--- a/include/openid.h	Tue Dec 28 19:57:25 2010 -0500
+++ b/include/openid.h	Wed Dec 29 12:16:32 2010 -0500
@@ -1,8 +1,12 @@
 #include <urweb/urweb.h>
 
+#include <openssl/dh.h>
+
 uw_unit uw_OpenidFfi_init(uw_context);
 
-typedef struct uw_OpenidFfi_discovery *uw_OpenidFfi_discovery;
+typedef struct {
+  uw_Basis_string endpoint, localId;
+} uw_OpenidFfi_discovery;
 
 uw_Basis_string uw_OpenidFfi_endpoint(uw_context, uw_OpenidFfi_discovery);
 uw_Basis_string uw_OpenidFfi_localId(uw_context, uw_OpenidFfi_discovery);
@@ -20,4 +24,14 @@
 uw_OpenidFfi_outputs uw_OpenidFfi_direct(uw_context, uw_Basis_string url, uw_OpenidFfi_inputs);
 uw_OpenidFfi_outputs uw_OpenidFfi_indirect(uw_context, uw_Basis_string fields);
 
+uw_Basis_string uw_OpenidFfi_sha1(uw_context, uw_Basis_string key, uw_Basis_string data);
 uw_Basis_string uw_OpenidFfi_sha256(uw_context, uw_Basis_string key, uw_Basis_string data);
+
+typedef DH *uw_OpenidFfi_dh;
+
+uw_Basis_string uw_OpenidFfi_modulus(uw_context, uw_OpenidFfi_dh);
+uw_Basis_string uw_OpenidFfi_generator(uw_context, uw_OpenidFfi_dh);
+uw_Basis_string uw_OpenidFfi_public(uw_context, uw_OpenidFfi_dh);
+
+uw_OpenidFfi_dh uw_OpenidFfi_generate(uw_context);
+uw_Basis_string uw_OpenidFfi_compute(uw_context, uw_OpenidFfi_dh, uw_Basis_string server_pub);

--- a/src/c/openid.c	Tue Dec 28 19:57:25 2010 -0500
+++ b/src/c/openid.c	Wed Dec 29 12:16:32 2010 -0500
@@ -5,6 +5,7 @@
 #include <openssl/buffer.h>
 #include <openssl/sha.h>
 #include <openssl/hmac.h>
+#include <openssl/dh.h>
 #include <curl/curl.h>
 #include <expat.h>
 
@@ -13,16 +14,15 @@
 #define BUF_MAX 10240
 #define BUF_INIT 1024
 
-struct uw_OpenidFfi_discovery {
-  uw_Basis_string endpoint, localId;
-};
+#define PRIME_LEN 64
+#define GENERATOR DH_GENERATOR_5
 
 uw_Basis_string uw_OpenidFfi_endpoint(uw_context ctx, uw_OpenidFfi_discovery d) {
-  return d->endpoint;
+  return d.endpoint;
 }
 
 uw_Basis_string uw_OpenidFfi_localId(uw_context ctx, uw_OpenidFfi_discovery d) {
-  return d->localId;
+  return d.localId;
 }
 
 uw_unit uw_OpenidFfi_init(uw_context ctx) {
@@ -45,7 +45,7 @@
 
 typedef struct {
   uw_context ctx;
-  uw_OpenidFfi_discovery d;
+  uw_OpenidFfi_discovery *d;
 } endpoint;
 
 static void XMLCALL startElement(void *userData, const XML_Char *name, const XML_Char **atts) {
@@ -91,7 +91,7 @@
   char *s;
   CURL *c = curl(ctx);
   curl_discovery_data cd = {};
-  uw_OpenidFfi_discovery dy = uw_malloc(ctx, sizeof(struct uw_OpenidFfi_discovery));
+  uw_OpenidFfi_discovery *dy = uw_malloc(ctx, sizeof(uw_OpenidFfi_discovery));
   endpoint ep = {ctx, dy};
   CURLcode code;
 
@@ -120,13 +120,10 @@
   code = curl_easy_perform(c);
   uw_pop_cleanup(ctx);
 
-  if (code || !ep.d->endpoint)
+  if (code || !dy->endpoint)
     return NULL;
-  else {
-    uw_OpenidFfi_discovery *dyp = malloc(sizeof(uw_OpenidFfi_discovery));
-    *dyp = ep.d;
-    return dyp;
-  }
+  else
+    return dy;
 }
 
 uw_OpenidFfi_inputs uw_OpenidFfi_createInputs(uw_context ctx) {
@@ -135,25 +132,28 @@
   return r;
 }
 
-static int okForPost(const char *s) {
-  for (; *s; ++s)
-    if (*s == '=' || *s == '&')
-      return 0;
-  return 1;
+static void postify(uw_OpenidFfi_inputs buf, uw_Basis_string s) {
+  for (; *s; ++s) {
+    switch (*s) {
+    case '=':
+      uw_buffer_append(buf, "%3D", 3);
+      break;
+    case '&':
+      uw_buffer_append(buf, "%26", 3);
+      break;
+    default:
+      uw_buffer_append(buf, s, 1);
+    }
+  }
 }
 
 uw_unit uw_OpenidFfi_addInput(uw_context ctx, uw_OpenidFfi_inputs buf, uw_Basis_string key, uw_Basis_string value) {
-  if (!okForPost(key))
-    uw_error(ctx, FATAL, "Invalid key for OpenID inputs");
-  if (!okForPost(value))
-    uw_error(ctx, FATAL, "Invalid value for OpenID inputs");
-
   if (uw_buffer_used(buf) > 0)
     uw_buffer_append(buf, "&", 1);
 
-  uw_buffer_append(buf, key, strlen(key));
+  postify(buf, key);
   uw_buffer_append(buf, "=", 1);
-  uw_buffer_append(buf, value, strlen(value));
+  postify(buf, value);
 
   return uw_unit_v;
 }
@@ -202,6 +202,8 @@
   } else {
     char *s;
 
+    printf("Result: %s\n", buf->start);
+
     s = buf->start;
     while (*s) {
       char *colon = strchr(s, ':'), *newline;
@@ -297,17 +299,30 @@
   return buff;
 }
 
-static void unbase64(unsigned char *input, int length, unsigned char *buffer, int bufferLength)
+static int unbase64(unsigned char *input, int length, unsigned char *buffer, int bufferLength)
 {
   BIO *b64, *bmem;
+  int n;
 
   b64 = BIO_new(BIO_f_base64());
   BIO_set_flags(b64, BIO_FLAGS_BASE64_NO_NL);
   bmem = BIO_new_mem_buf(input, length);
   BIO_push(b64, bmem);
-  BIO_read(b64, buffer, bufferLength);
+  n = BIO_read(b64, buffer, bufferLength);
 
   BIO_free_all(bmem);
+
+  return n;
+}
+
+uw_Basis_string uw_OpenidFfi_sha1(uw_context ctx, uw_Basis_string key, uw_Basis_string data) {
+  unsigned char keyBin[SHA_DIGEST_LENGTH], out[EVP_MAX_MD_SIZE];
+  unsigned outLen;
+
+  unbase64((unsigned char *)key, strlen(key), keyBin, sizeof keyBin);
+
+  HMAC(EVP_sha1(), keyBin, sizeof keyBin, (unsigned char *)data, strlen(data), out, &outLen);
+  return base64(ctx, out, outLen);
 }
 
 uw_Basis_string uw_OpenidFfi_sha256(uw_context ctx, uw_Basis_string key, uw_Basis_string data) {
@@ -315,8 +330,81 @@
   unsigned outLen;
 
   unbase64((unsigned char *)key, strlen(key), keyBin, sizeof keyBin);
-  memset(key, sizeof key, 0);
 
   HMAC(EVP_sha256(), keyBin, sizeof keyBin, (unsigned char *)data, strlen(data), out, &outLen);
   return base64(ctx, out, outLen);
 }
+
+static uw_Basis_string btwoc(uw_context ctx, const BIGNUM *n) {
+  int len = BN_num_bytes(n), i;
+  unsigned char bytes[len+1];
+
+  bytes[0] = 0;
+  BN_bn2bin(n, bytes+1);
+
+  for (i = 1; i <= len; ++i)
+    if (bytes[i]) {
+      if (bytes[i] & 0x80)
+        --i;
+      break;
+    }
+
+  if (i > len)
+    i = len;
+
+  return base64(ctx, bytes+i, len+1-i);
+}
+
+static BIGNUM *unbtwoc(uw_context ctx, uw_Basis_string s) {
+  unsigned char bytes[1024];
+  int len;
+
+  len = unbase64((unsigned char *)s, strlen(s), bytes, sizeof bytes);
+  return BN_bin2bn(bytes, len, NULL);
+}
+
+uw_Basis_string uw_OpenidFfi_modulus(uw_context ctx, uw_OpenidFfi_dh dh) {
+  return btwoc(ctx, dh->p);
+}
+
+uw_Basis_string uw_OpenidFfi_generator(uw_context ctx, uw_OpenidFfi_dh dh) {
+  return btwoc(ctx, dh->g);
+}
+
+uw_Basis_string uw_OpenidFfi_public(uw_context ctx, uw_OpenidFfi_dh dh) {
+  return btwoc(ctx, dh->pub_key);
+}
+
+static void free_DH(void *data, int will_retry) {
+  DH *dh = data;
+  DH_free(dh);
+}
+
+uw_OpenidFfi_dh uw_OpenidFfi_generate(uw_context ctx) {
+  DH *dh = DH_new();
+
+  uw_register_transactional(ctx, dh, NULL, NULL, free_DH);
+
+  DH_generate_parameters_ex(dh, PRIME_LEN, GENERATOR, NULL);
+
+  if (DH_generate_key(dh) != 1)
+    uw_error(ctx, FATAL, "Diffie-Hellman key generation failed");
+
+  return dh;
+}
+
+uw_Basis_string uw_OpenidFfi_compute(uw_context ctx, uw_OpenidFfi_dh dh, uw_Basis_string server_pub) {
+  BIGNUM *bn = unbtwoc(ctx, server_pub);
+  unsigned char secret[DH_size(dh)];
+  int size;
+
+  uw_push_cleanup(ctx, (void (*)(void *))BN_free, bn);
+
+  size = DH_compute_key(secret, bn, dh);
+  if (size == -1)
+    uw_error(ctx, FATAL, "Diffie-Hellman key computation failed");
+
+  uw_pop_cleanup(ctx);
+
+  return base64(ctx, secret, size);
+}

--- a/src/ur/lib.urp	Tue Dec 28 19:57:25 2010 -0500
+++ b/src/ur/lib.urp	Wed Dec 29 12:16:32 2010 -0500
@@ -5,7 +5,10 @@
 effectful OpenidFfi.discover
 effectful OpenidFfi.createInputs
 effectful OpenidFfi.addInput
+effectful OpenidFfi.direct
 effectful OpenidFfi.indirect
+effectful OpenidFfi.generate
+effectful OpenidFfi.compute
 
 $/string
 $/option

--- a/src/ur/openid.ur	Tue Dec 28 19:57:25 2010 -0500
+++ b/src/ur/openid.ur	Wed Dec 29 12:16:32 2010 -0500
@@ -27,48 +27,164 @@
     OpenidFfi.addInput is "openid.ns" "http://specs.openid.net/auth/2.0";
     return is
 
-table associations : { Endpoint : string, Handle : string, Key : string, Expires : time }
+datatype association_type = HMAC_SHA1 | HMAC_SHA256
+datatype association_session_type = NoEncryption | DH_SHA1 | DH_SHA256
+
+table associations : { Endpoint : string, Handle : string, Typ : serialized association_type, Key : string, Expires : time }
   PRIMARY KEY Endpoint
 
-datatype association = Association of {Handle : string, Key : string} | AssError of string
+datatype association = Association of {Handle : string, Typ : association_type, Key : string}
+                     | AssError of string
+                     | AssAlternate of {Atype : association_type, Stype : association_session_type}
 
-fun association url =
-    secret <- oneOrNoRows1 (SELECT associations.Handle, associations.Key
+fun atype_show v =
+    case v of
+        HMAC_SHA1 => "HMAC-SHA1"
+      | HMAC_SHA256 => "HMAC-SHA256"
+
+val show_atype = mkShow atype_show
+
+fun stype_show v =
+    case v of
+        NoEncryption => "no-encryption"
+      | DH_SHA1 => "DH-SHA1"
+      | DH_SHA256 => "DH-SHA256"
+
+val show_stype = mkShow stype_show
+
+fun atype_read s =
+    case s of
+        "HMAC-SHA1" => Some HMAC_SHA1
+      | "HMAC-SHA256" => Some HMAC_SHA256
+      | _ => None
+
+val read_atype = mkRead' atype_read "association type"
+
+fun stype_read s =
+    case s of
+        "no-encryption" => Some NoEncryption
+      | "DH-SHA1" => Some DH_SHA1
+      | "DH-SHA256" => Some DH_SHA256
+      | _ => None
+
+val read_stype = mkRead' stype_read "association session type"
+
+fun atype_eq v1 v2 =
+    case (v1, v2) of
+        (HMAC_SHA1, HMAC_SHA1) => True
+      | (HMAC_SHA256, HMAC_SHA256) => True
+      | _ => False
+
+val eq_atype = mkEq atype_eq
+
+fun stype_eq v1 v2 =
+    case (v1, v2) of
+        (NoEncryption, NoEncryption) => True
+      | (DH_SHA1, DH_SHA1) => True
+      | (DH_SHA256, DH_SHA256) => True
+      | _ => False
+
+val eq_stype = mkEq stype_eq
+
+fun errorResult atype stype os =
+    case OpenidFfi.getOutput os "error" of
+        Some v =>
+        (case (OpenidFfi.getOutput os "error_code", OpenidFfi.getOutput os "assoc_type", OpenidFfi.getOutput os "session_type") of
+             (Some "unsupported-type", at, st) => Some (AssAlternate {Atype = Option.get atype (Option.bind read at),
+                                                                      Stype = Option.get stype (Option.bind read st)})
+           | _ => Some (AssError ("OP error during association: " ^ v)))
+      | None => None
+
+fun associateNoEncryption url atype =
+    is <- createInputs;
+    OpenidFfi.addInput is "openid.mode" "associate";
+    OpenidFfi.addInput is "openid.assoc_type" (show atype);
+    OpenidFfi.addInput is "openid.session_type" (show NoEncryption);
+
+    os <- OpenidFfi.direct url is;
+    case errorResult atype NoEncryption os of
+        Some v => return v
+      | None =>
+        case (OpenidFfi.getOutput os "assoc_handle", OpenidFfi.getOutput os "mac_key", OpenidFfi.getOutput os "expires_in") of
+            (Some handle, Some key, Some expires) =>
+            (case read expires of
+                 None => return (AssError "Invalid 'expires_in' field")
+               | Some expires =>
+                 tm <- now;
+                 dml (INSERT INTO associations (Endpoint, Handle, Typ, Key, Expires)
+                      VALUES ({[url]}, {[handle]}, {[serialize atype]}, {[key]}, {[addSeconds tm expires]}));
+                 return (Association {Handle = handle, Typ = atype, Key = key}))
+          | (None, _, _) => return (AssError "Missing assoc_handle")
+          | (_, None, _) => return (AssError "Missing mac_key")
+          | _ => return (AssError "Missing expires_in")
+
+fun associateDh url atype stype =
+    dh <- OpenidFfi.generate;
+
+    is <- createInputs;
+    OpenidFfi.addInput is "openid.mode" "associate";
+    OpenidFfi.addInput is "openid.assoc_type" (show atype);
+    OpenidFfi.addInput is "openid.session_type" (show stype);
+    OpenidFfi.addInput is "openid.dh_modulus" (OpenidFfi.modulus dh);
+    OpenidFfi.addInput is "openid.dh_gen" (OpenidFfi.generator dh);
+    OpenidFfi.addInput is "openid.dh_consumer_public" (OpenidFfi.public dh);
+
+    os <- OpenidFfi.direct url is;
+    case errorResult atype stype os of
+        Some v => return v
+      | None =>
+        case (OpenidFfi.getOutput os "assoc_handle", OpenidFfi.getOutput os "dh_server_public",
+              OpenidFfi.getOutput os "enc_mac_key", OpenidFfi.getOutput os "expires_in") of
+                (Some handle, Some pub, Some mac, Some expires) =>
+                (case read expires of
+                     None => return (AssError "Invalid 'expires_in' field")
+                   | Some expires =>
+                     key <- OpenidFfi.compute dh pub;
+                     tm <- now;
+                     dml (INSERT INTO associations (Endpoint, Handle, Typ, Key, Expires)
+                          VALUES ({[url]}, {[handle]}, {[serialize atype]}, {[key]}, {[addSeconds tm expires]}));
+                     return (Association {Handle = handle, Typ = atype, Key = key}))
+              | (None, _, _, _) => return (AssError "Missing assoc_handle")
+              | (_, None, _, _) => return (AssError "Missing dh_server_public")
+              | (_, _, None, _) => return (AssError "Missing enc_mac_key")
+              | _ => return (AssError "Missing expires_in")
+
+fun oldAssociation url =
+    secret <- oneOrNoRows1 (SELECT associations.Handle, associations.Typ, associations.Key
                             FROM associations
                             WHERE associations.Endpoint = {[url]});
     case secret of
+        Some r => return (Some (r -- #Typ ++ {Typ = deserialize r.Typ}))
+      | None => return None
+
+fun newAssociation url atype stype =
+    case stype of
+        NoEncryption => associateNoEncryption url atype
+      | _ => associateDh url atype stype
+
+fun association atype stype url =
+    secret <- oldAssociation url;
+    case secret of
         Some r => return (Association r)
       | None =>
-        is <- createInputs;
-        OpenidFfi.addInput is "openid.mode" "associate";
-        OpenidFfi.addInput is "openid.assoc_type" "HMAC-SHA256";
-        OpenidFfi.addInput is "openid.session_type" "no-encryption";
-
-        debug ("Contacting " ^ url);
-
-        os <- OpenidFfi.direct url is;
-        case OpenidFfi.getOutput os "error" of
-            Some v => return (AssError v)
-          | None =>
-            case (OpenidFfi.getOutput os "assoc_handle", OpenidFfi.getOutput os "mac_key", OpenidFfi.getOutput os "expires_in") of
-                (Some handle, Some key, Some expires) =>
-                (case read expires of
-                     None => return (AssError "Invalid 'expires_in' field")
-                   | Some expires =>
-                     tm <- now;
-                     dml (INSERT INTO associations (Endpoint, Handle, Key, Expires)
-                          VALUES ({[url]}, {[handle]}, {[key]}, {[addSeconds tm expires]}));
-                     return (Association {Handle = handle, Key = key}))
-              | (None, _, _) => return (AssError "Missing assoc_handle")
-              | (_, None, _) => return (AssError "Missing mac_key")
-              | _ => return (AssError "Missing fields in response from OP")
+        stype <- return (case (stype, String.isPrefix {Full = url, Prefix = "https://"}) of
+                             (NoEncryption, False) => DH_SHA256
+                           | _ => stype);
+        r <- newAssociation url atype stype;
+        case r of
+            AssAlternate alt =>
+            if alt.Atype = atype && alt.Stype = stype then
+                return (AssError "Suggested new modes match old ones!")
+            else
+                newAssociation url alt.Atype alt.Stype
+          | v => return v
 
 fun eatFragment s =
     case String.split s #"#" of
         Some (_, s') => s'
       | _ => s
 
-datatype handle_result = HandleOk of {Endpoint : string, Key : string} | HandleError of string
+datatype handle_result = HandleOk of {Endpoint : string, Typ : association_type, Key : string} | HandleError of string
 
 fun verifyHandle os id =
     ep <- discover (eatFragment id);
@@ -78,14 +194,14 @@
         case OpenidFfi.getOutput os "openid.assoc_handle" of
             None => return (HandleError "Missing association handle in response")
           | Some handle =>
-            assoc <- association ep;
+            assoc <- oldAssociation ep;
             case assoc of
-                AssError s => return (HandleError s)
-              | Association assoc =>
+                None => return (HandleError "Couldn't find association handle")
+              | Some assoc =>
                 if assoc.Handle <> handle then
                     return (HandleError "Association handles don't match")
                 else
-                    return (HandleOk {Endpoint = ep, Key = assoc.Key})
+                    return (HandleOk {Endpoint = ep, Typ = assoc.Typ, Key = assoc.Key})
 
 table nonces : { Endpoint : string, Nonce : string, Expires : time }
   PRIMARY KEY (Endpoint, Nonce)
@@ -123,7 +239,7 @@
                          VALUES ({[ep]}, {[nonce]}, {[exp]}));
                     return None
 
-fun verifySig os key =
+fun verifySig os atype key =
     case OpenidFfi.getOutput os "openid.signed" of
         None => return (Some "Missing openid.signed in OP response")
       | Some signed =>
@@ -153,7 +269,9 @@
                     None => return (Some "openid.signed mentions missing field")
                   | Some nvps =>
                     let
-                        val sign' = OpenidFfi.sha256 key nvps
+                        val sign' = case atype of
+                                        HMAC_SHA256 => OpenidFfi.sha256 key nvps
+                                      | HMAC_SHA1 => OpenidFfi.sha1 key nvps
                     in
                         debug ("Fields: " ^ signed);
                         debug ("Nvps: " ^ nvps);
@@ -187,7 +305,7 @@
                          errO <- verifyHandle os id;
                          case errO of
                              HandleError s => error <xml>{[s]}</xml>
-                           | HandleOk {Endpoint = ep, Key = key} =>
+                           | HandleOk {Endpoint = ep, Typ = atype, Key = key} =>
                              errO <- verifyReturnTo os;
                              case errO of
                                  Some s => error <xml>{[s]}</xml>
@@ -196,7 +314,7 @@
                                  case errO of
                                      Some s => error <xml>{[s]}</xml>
                                    | None =>
-                                     errO <- verifySig os key;
+                                     errO <- verifySig os atype key;
                                      case errO of
                                          Some s => error <xml>{[s]}</xml>
                                        | None => return <xml>Identity: {[id]}</xml>)
@@ -211,14 +329,15 @@
         else
             return None
 
-fun authenticate id =
+fun authenticate atype stype id =
     dy <- discover id;
     case dy of
         None => return "Discovery failed"
       | Some dy =>
-        assoc <- association dy;
+        assoc <- association atype stype dy;
         case assoc of
-            AssError msg => return msg
+            AssError msg => return ("Association failure: " ^ msg)
+          | AssAlternate _ => return "Association failure: server didn't accept its own alternate association modes"
           | Association assoc =>
             redirect (bless (dy ^ "?openid.ns=http://specs.openid.net/auth/2.0&openid.mode=checkid_setup&openid.claimed_id="
                              ^ id ^ "&openid.identity=http://specs.openid.net/auth/2.0/identifier_select&openid.assoc_handle="

--- a/src/ur/openid.urs	Tue Dec 28 19:57:25 2010 -0500
+++ b/src/ur/openid.urs	Wed Dec 29 12:16:32 2010 -0500
@@ -1,3 +1,6 @@
-val authenticate : string -> transaction string
+datatype association_type = HMAC_SHA1 | HMAC_SHA256
+datatype association_session_type = NoEncryption | DH_SHA1 | DH_SHA256
+
+val authenticate : association_type -> association_session_type -> string -> transaction string
 (* Doesn't return normally if everything goes as planned.
  * Instead, the user is redirected to his OP to authenticate there. *)

--- a/src/ur/openidFfi.urs	Tue Dec 28 19:57:25 2010 -0500
+++ b/src/ur/openidFfi.urs	Wed Dec 29 12:16:32 2010 -0500
@@ -15,4 +15,13 @@
 val direct : string -> inputs -> transaction outputs
 val indirect : queryString -> transaction outputs
 
+val sha1 : string -> string -> string
 val sha256 : string -> string -> string
+
+type dh
+val modulus : dh -> string
+val generator: dh -> string
+val public : dh -> string
+
+val generate : transaction dh
+val compute : dh -> string -> transaction string

--- a/tests/test.ur	Tue Dec 28 19:57:25 2010 -0500
+++ b/tests/test.ur	Wed Dec 29 12:16:32 2010 -0500
@@ -1,5 +1,5 @@
 fun auth r =
-    msg <- Openid.authenticate r.Id;
+    msg <- Openid.authenticate Openid.HMAC_SHA256 Openid.NoEncryption r.Id;
     error <xml>{[msg]}</xml>
 
 fun main () = return <xml><body>