3 files changed, 27 insertions, 20 deletions

diff --git a/Auth/OpenID/SReg.php b/Auth/OpenID/SReg.php
index 287b8d8..208a6ab 100644
--- a/Auth/OpenID/SReg.php
+++ b/Auth/OpenID/SReg.php
@@ -27,7 +27,7 @@
  *   //   the fields in sreg_response were requested ]
  *   $sreg_resp = Auth_OpenID_SRegResponse::extractResponse(
  *                                  $sreg_req, $user_data);
- *   $sreg_resp->addToOpenIDResponse($openid_response);
+ *   $sreg_resp->toMessage($openid_response->fields);
  *
  * 3. The relying party uses {@link
  * Auth_OpenID_SRegResponse::fromSuccessResponse} to extract the data
@@ -465,7 +465,7 @@ class Auth_OpenID_SRegResponse extends Auth_OpenID_SRegBase {
      * Returns a simple registration response containing the data that
      * was supplied with the C{id_res} response.
      */
-    function fromSuccessResponse($success_response, $signed_only=true)
+    function fromSuccessResponse(&$success_response, $signed_only=true)
     {
         global $Auth_OpenID_sreg_data_fields;
 
@@ -491,18 +491,9 @@ class Auth_OpenID_SRegResponse extends Auth_OpenID_SRegBase {
         return $obj;
     }
 
-    /**
-     * Add the data fields contained in this simple registration
-     * response to the supplied message, in the appropriate namespace.
-     *
-     * response_message: The OpenID id_res response message that will
-     * be returned to the relying party
-     *
-     * Returns nothing; updates the response_message
-     */
-    function addToOpenIDResponse(&$response_message)
+    function getExtensionArgs()
     {
-        $response_message->updateArgs($this->ns_uri, $this->data);
+        return $this->data;
     }
 
     // Read-only dictionary interface
@@ -544,7 +535,7 @@ function Auth_OpenID_sendSRegFields(&$openid_request, $data, &$openid_response)
                                                    $openid_request->message);
     $sreg_response = Auth_OpenID_SRegResponse::extractResponse(
                                                    $sreg_request, $data);
-    $sreg_response->addToOpenIDResponse($openid_response->fields);
+    $sreg_response->toMessage($openid_response->fields);
 }
 
 ?>
\ No newline at end of file
diff --git a/Tests/Auth/OpenID/SReg.php b/Tests/Auth/OpenID/SReg.php
index f807aa6..395daf1 100644
--- a/Tests/Auth/OpenID/SReg.php
+++ b/Tests/Auth/OpenID/SReg.php
@@ -576,23 +576,23 @@ class SRegResponseTest extends PHPUnit_TestCase {
 }
 
 class SendFieldsTest extends PHPUnit_TestCase {
-    function test()
+    function _test($uri)
     {
         // Create a request message with simple registration fields
         $sreg_req = Auth_OpenID_SRegRequest::build(array('nickname', 'email'),
                                                    array('fullname'));
-        $req_msg = new Auth_OpenID_Message();
+        $req_msg = new Auth_OpenID_Message($uri);
         $req_msg->updateArgs(Auth_OpenID_SREG_NS_URI,
                              $sreg_req->getExtensionArgs());
 
         $req = new Auth_OpenID_Request();
-        $req->message = $req_msg;
+        $req->message =& $req_msg;
         $req->namespace = $req_msg->getOpenIDNamespace();
 
         // -> send checkid_* request
 
         // Create an empty response message
-        $resp_msg = new Auth_OpenID_Message();
+        $resp_msg = new Auth_OpenID_Message($uri);
         $resp = new Auth_OpenID_ServerResponse($req);
         $resp->fields = $resp_msg;
 
@@ -619,6 +619,14 @@ class SendFieldsTest extends PHPUnit_TestCase {
                                   'fullname' => 'Leonhard Euler'),
                             $sreg_data_resp);
     }
+
+    function test()
+    {
+        foreach (array(Auth_OpenID_OPENID1_NS,
+                       Auth_OpenID_OPENID2_NS) as $uri) {
+            $this->_test($uri);
+        }
+    }
 }
 
 class Tests_Auth_OpenID_SReg extends PHPUnit_TestSuite {
diff --git a/examples/server/lib/common.php b/examples/server/lib/common.php
index badda5f..600b987 100644
--- a/examples/server/lib/common.php
+++ b/examples/server/lib/common.php
@@ -66,9 +66,17 @@ function doAuth($info, $trusted=null, $fail_cancels=false,
                            'language' => 'eu',
                            'timezone' => 'America/New_York');
 
-        Auth_OpenID_sendSRegFields($info, $sreg_data,
-                                   $response);
+        // Add the simple registration response values to the OpenID
+        // response message.
+        $sreg_request = Auth_OpenID_SRegRequest::fromOpenIDRequest(
+                                              $info->message);
 
+        $sreg_response = Auth_OpenID_SRegResponse::extractResponse(
+                                              $sreg_request, $sreg_data);
+
+        $sreg_response->toMessage($response->fields);
+
+        // Generate a response to send to the user agent.
         $webresponse =& $server->encodeResponse($response);
 
         $new_headers = array();