Use TOTP for magic codes

server/resources/migrations/100_totp_magic_code.down.sql

@@ -0,0 +1,2 @@
+alter table apps drop column totp_secret_key_enc;
+alter table apps drop column totp_expiry_minutes;

server/resources/migrations/100_totp_magic_code.up.sql

@@ -0,0 +1,2 @@
+alter table apps add column totp_secret_key_enc bytea;
+alter table apps add column totp_expiry_minutes integer;

server/src/instant/admin/routes.clj

@@ -500,7 +500,7 @@
 (defn magic-code-post [req]
   (let [{:keys [app-id]} (req->app-id-authed! req :data/write)
         email (ex/get-param! req [:body :email] email/coerce)
-        {:keys [code]} (app-user-magic-code-model/create! {:app-id app-id :email email})]
+        code (app-user-magic-code-model/create! {:app-id app-id :email email})]
     (response/ok {:code code})))
 
 (defn send-magic-code-post [req]

server/src/instant/flags.clj

@@ -245,15 +245,41 @@
 (defn test-rule-wheres? []
   (:rule-where-testing (query-result)))
 
+(def ^:dynamic *toggle-overrides* nil)
+
 (defn toggled?
   ([key]
-   (get-in (query-result) [:toggles key]))
+   (if *toggle-overrides*
+     (-> (query-result)
+         (get :toggles)
+         (merge *toggle-overrides*)
+         (get key))
+     (get-in (query-result) [:toggles key])))
   ([key not-found]
-   (get-in (query-result) [:toggles key] not-found)))
+   (if *toggle-overrides*
+     (-> (query-result)
+         (get :toggles)
+         (merge *toggle-overrides*)
+         (get key not-found))
+     (get-in (query-result) [:toggles key] not-found))))
+
+(def ^:dynamic *flag-overrides* nil)
 
 (defn flag
-  ([key] (get-in (query-result) [:flags key]))
-  ([key not-found] (get-in (query-result) [:flags key] not-found)))
+  ([key]
+   (if *flag-overrides*
+     (-> (query-result)
+         (get :flags)
+         (merge *flag-overrides*)
+         (get key))
+     (get-in (query-result) [:flags key])))
+  ([key not-found]
+   (if *flag-overrides*
+     (-> (query-result)
+         (get :flags)
+         (merge *flag-overrides*)
+         (get key not-found))
+     (get-in (query-result) [:flags key] not-found))))
 
 (defn handle-receive-timeout [app-id]
   (get-in (query-result) [:handle-receive-timeout app-id]))
@@ -334,3 +360,33 @@
    Set to 0 to disable."
   []
   (flag :magic-code-rate-limit-per-hour 20))
+
+;; TOTP flags
+;; Process for rolling out TOTP
+;; 1. Deploy the code and wait for 24 hours
+;; 2. Flip the validate-with-totp? toggle to true
+;; 3. Wait for a while to make sure everything looks good
+;; 4. Flip the dual-write-totp toggle to false
+;; 5. Remove the code that writes to the magic codes table
+;; 6. Later--remove the magic codes completely
+
+(defn generate-with-totp?
+  "Returns true if we should generate codes with totp. Need to run this for 24 hours before
+   turning on validate-with-totp"
+  []
+  (toggled? :generate-with-totp true))
+
+(defn validate-with-totp?
+  "Returns true if we should validate codes with totp. Need to generate-with-totp for 24 hours before
+   turning on validate-with-totp"
+  []
+  (toggled? :validate-with-totp false))
+
+(defn dual-write-totp?
+  "Returns true if we should write our totp codes to the magic codes table."
+  []
+  (toggled? :dual-write-totp true))
+
+(defn totp-default-expiry-seconds
+  []
+  (flag :totp-default-expiry-seconds 86400))

server/src/instant/model/app.clj

@@ -346,6 +346,35 @@
                                  WHERE a.id = ?::uuid"
                                new-creator-id id])))))
 
+(defn ensure-totp-secret!
+  ([params]
+   (ensure-totp-secret! (aurora/conn-pool :write) params))
+  ([conn {:keys [id]}]
+   (let [secret-key (crypt-util/random-bytes 32)
+         enc-secret-key (crypt-util/aead-encrypt {:plaintext secret-key
+                                                  :associated-data (uuid-util/->bytes id)})]
+     (with-cache-invalidation id
+       (sql/execute-one! ::ensure-totp-secret!
+                         conn
+                         (hsql/format {:update :apps
+                                       :set {:totp-secret-key-enc enc-secret-key}
+                                       :where [:and
+                                               [:= :id id]
+                                               [:= :totp-secret-key-enc nil]]}))))))
+
+(defn get-totp-secret-key
+  ([params]
+   (get-totp-secret-key (aurora/conn-pool :read)
+                        (aurora/conn-pool :write)
+                        params))
+  ([read-conn write-conn {:keys [id]}]
+   (let [secret-key-enc (or (:totp_secret_key_enc (get-by-id! read-conn {:id id}))
+                            (:totp_secret_key_enc (ensure-totp-secret! write-conn {:id id}))
+                            ;; 2nd read in case two txes tried to initialize the key at the same time
+                            (:totp_secret_key_enc (get-by-id! read-conn {:id id})))]
+     (crypt-util/aead-decrypt {:ciphertext secret-key-enc
+                               :associated-data (uuid-util/->bytes id)}))))
+
 (defn clear-by-id!
   "Deletes attrs, rules, and triples for the specified app_id"
   ([params] (clear-by-id! (aurora/conn-pool :write) params))

server/src/instant/model/app_user_magic_code.clj

@@ -1,10 +1,12 @@
 (ns instant.model.app-user-magic-code
   (:require
+   [instant.flags :as flags]
    [instant.jdbc.aurora :as aurora]
    [instant.model.app :as app-model]
    [instant.model.app-user :as app-user-model]
    [instant.model.instant-user :as instant-user-model]
    [instant.system-catalog-ops :refer [update-op]]
+   [instant.totp :as totp]
    [instant.util.crypt :as crypt-util]
    [instant.util.exception :as ex]
    [instant.util.string :refer [rand-num-str]])
@@ -24,47 +26,81 @@
     (let [created-at ^Date (:created_at magic-code)]
       (< (+ (.getTime created-at) ttl-ms) (System/currentTimeMillis)))))
 
+(defn totp-secret-key [app-id ^String email]
+  (let [app-secret-key (app-model/get-totp-secret-key {:id app-id})
+        derived-key (crypt-util/hmac-256 app-secret-key (.getBytes email))]
+    derived-key))
+
+(defn generate-totp [app-id ^String email]
+  (let [secret-key (totp-secret-key app-id email)]
+    (totp/generate-totp secret-key)))
+
 (defn create!
   ([params]
    (create! (aurora/conn-pool :write) params))
-  ([conn {:keys [app-id email id code]}]
+  ([conn {:keys [app-id email id]}]
    (let [id   (or id (random-uuid))
-         code (or code (rand-code))]
-     (update-op
-      conn
-      {:app-id app-id
-       :etype etype}
-      (fn [{:keys [transact! get-entity]}]
-        (transact! [{:id       id
-                     :etype    etype
-                     :codeHash (-> code
-                                   crypt-util/str->sha256
-                                   crypt-util/bytes->hex-string)
-                     :email    email}])
-        (assoc (get-entity id)
-               :code code))))))
+         code (if (flags/generate-with-totp?)
+                (generate-totp app-id email)
+                (rand-code))]
+
+     (when (or (not (flags/validate-with-totp?))
+               (not (flags/generate-with-totp?))
+               (flags/dual-write-totp?))
+       (update-op
+        conn
+        {:app-id app-id
+         :etype etype}
+        (fn [{:keys [transact!]}]
+          (transact! [{:id       id
+                       :etype    etype
+                       :codeHash (-> code
+                                     crypt-util/str->sha256
+                                     crypt-util/bytes->hex-string)
+                       :email    email}]))))
+     code)))
+
+
+(defn validate-totp! [app-id ^String email ^String code]
+  (let [secret-key (totp-secret-key app-id email)
+        expiry-periods (or (some-> (app-model/get-by-id! {:id app-id})
+                                   :totp_expiry_minutes
+                                   (max 5) ;; Minimum of 5 minutes
+                                   (min 1440) ;; Maximum of 1 day
+                                   (* 60)
+                                   (/ totp/default-time-step)
+                                   (Math/ceil))
+                           (/ (flags/totp-default-expiry-seconds) totp/default-time-step))]
+    ;; Have to add 1 extra period in case the code was generated near the
+    ;; end of a period
+    (when-not (totp/valid-totp? secret-key (inc expiry-periods) code)
+      (ex/throw-expiration-err! :app-user-magic-code {:args [{:code code
+                                                              :email email}]}))))
 
 (defn consume!
   ([params]
    (consume! (aurora/conn-pool :write) params))
   ([conn {:keys [email code app-id] :as params}]
-   (update-op
-    conn
-    {:app-id app-id
-     :etype etype}
-    (fn [{:keys [get-entity-where delete-entity!]}]
-      (let [code-hash (-> code
-                          crypt-util/str->sha256
-                          crypt-util/bytes->hex-string)
-            {code-id :id} (get-entity-where
-                           {:codeHash code-hash
-                            :email    email})]
-        (ex/assert-record! code-id :app-user-magic-code {:args [params]})
-        (let [code (delete-entity! code-id)]
-          (ex/assert-record! code :app-user-magic-code {:args [params]})
-          (when (expired? code)
-            (ex/throw-expiration-err! :app-user-magic-code {:args [params]}))
-          code))))))
+   (when (or (not (flags/validate-with-totp?))
+             (flags/dual-write-totp?))
+     (update-op
+      conn
+      {:app-id app-id
+       :etype etype}
+      (fn [{:keys [get-entity-where delete-entity!]}]
+        (let [code-hash (-> code
+                            crypt-util/str->sha256
+                            crypt-util/bytes->hex-string)
+              {code-id :id} (get-entity-where
+                             {:codeHash code-hash
+                              :email    email})]
+          (ex/assert-record! code-id :app-user-magic-code {:args [params]})
+          (let [code (delete-entity! code-id)]
+            (ex/assert-record! code :app-user-magic-code {:args [params]})
+            (when (expired? code)
+              (ex/throw-expiration-err! :app-user-magic-code {:args [params]})))))))
+   (when (flags/validate-with-totp?)
+     (validate-totp! app-id email code))))
 
 (comment
   (def instant-user (instant-user-model/get-by-email

server/src/instant/runtime/magic_code_auth.clj

@@ -83,7 +83,7 @@
 (defn send! [{:keys [app-id email] :as req}]
   (check-rate-limit! req)
   (let [app             (app-model/get-by-id! {:id app-id})
-        {:keys [code]}  (app-user-magic-code-model/create! (select-keys req [:app-id :email]))
+        code            (app-user-magic-code-model/create! (select-keys req [:app-id :email]))
         template        (app-email-template-model/get-by-app-id-and-email-type
                          {:app-id app-id
                           :email-type "magic-code"})
@@ -178,10 +178,10 @@
   (def app (first (app-model/get-all-for-user {:user-id (:id instant-user)})))
   (def runtime-user (app-user-model/get-by-email {:app-id (:id app)
                                                   :email "stopa@instantdb.com"}))
-  (def m
-    (:magic-code (app-user-magic-code-model/create! {:app-id (:id app) :email "stopa@instantdb.com"})))
+  (def code
+    (app-user-magic-code-model/create! {:app-id (:id app) :email "stopa@instantdb.com"}))
 
   (verify! {:app-id (:id app) :email "stopa@instantdb.com" :code "0"})
 
-  (verify! {:app-id (:id app) :email "stopa@instantdb.com" :code (:code m)}))
+  (verify! {:app-id (:id app) :email "stopa@instantdb.com" :code code}))
 

server/src/instant/runtime/routes.clj

@@ -117,12 +117,12 @@
                                                   :email "stopa@instantdb.com"}))
 
   (def m (app-user-magic-code-model/create!
-          {:id (random-uuid) :user-id (:id runtime-user) :code (app-user-magic-code-model/rand-code)
+          {:id (random-uuid) :user-id (:id runtime-user)
            :app-id (:id app)}))
-  (verify-magic-code-post {:body {:email "stopainstantdb.com" :code (:code m)}})
-  (verify-magic-code-post {:body {:email "stopa@instantdb.com" :code (:code m)}})
+  (verify-magic-code-post {:body {:email "stopainstantdb.com" :code m}})
+  (verify-magic-code-post {:body {:email "stopa@instantdb.com" :code m}})
   (verify-magic-code-post {:body {:email "stopa@instantdb.com" :code "0" :app-id (:id app)}})
-  (verify-magic-code-post {:body {:email "stopa@instantdb.com" :code (:code m) :app-id (:id app)}}))
+  (verify-magic-code-post {:body {:email "stopa@instantdb.com" :code m :app-id (:id app)}}))
 
 ;; -----
 ;; Guest sign in

server/src/instant/totp.clj

@@ -0,0 +1,83 @@
+(ns instant.totp
+  (:require
+   [instant.util.crypt :as crypt-util]
+   [clojure.string])
+  (:import
+   (java.time Instant)))
+
+;; Time step is 5 minutes (in seconds), which half the default expiration. If an
+;; app accepts tokens for longer, we just check all of the previous 5 minute
+;; intervals until we find a matching token or exceed the max time. In the worst
+;; case (24 hours), we'll have to do 288 checks.
+;; We can't change this without adding extra code for backwards compatibility
+(def default-time-step 300)
+
+(def digit-count 6)
+
+;; TOTP generation follows the reference impl in https://www.rfc-editor.org/rfc/rfc6238#page-9
+
+(defn left-pad [s len]
+  (let [pad-len (- len (count s))]
+    (if (pos? pad-len)
+      (str (clojure.string/join (repeat pad-len "0")) s)
+      s)))
+
+(def digits-power
+  [1
+   10
+   100
+   1000
+   10000
+   100000
+   1000000
+   10000000
+   100000000])
+
+;; For use in testing
+(def ^:dynamic *now* nil)
+
+(defn generate-totp
+  "Generates a TOTP code. For testing, it accepts the number of
+   digits and a time step, but you should always use the default
+   values in production."
+  ([^bytes secret-key]
+   (generate-totp secret-key (or *now* (Instant/now))))
+  ([^bytes secret-key ^Instant time]
+   (generate-totp secret-key time 6 default-time-step))
+  ([^bytes secret-key ^Instant time code-digits time-step]
+   (let [t (/ (.getEpochSecond time)
+              time-step)
+         t-bytes (-> t
+                     (Long/toHexString)
+                     (.toUpperCase)
+                     (left-pad 16)
+                     (crypt-util/hex-string->bytes))
+         hash (crypt-util/hmac-256 secret-key t-bytes)
+         offset (bit-and (aget hash (dec (alength hash)))
+                         0xf)
+         binary (bit-or
+                 (bit-shift-left (bit-and (aget hash offset) 0x7f)
+                                 24)
+                 (bit-shift-left (bit-and (aget hash (+ offset 1)) 0xff)
+                                 16)
+                 (bit-shift-left (bit-and (aget hash (+ offset 2)) 0xff)
+                                 8)
+                 (bit-and (aget hash (+ offset 3)) 0xff))
+         otp (mod binary (nth digits-power code-digits))]
+     (left-pad (Integer/toString otp) code-digits))))
+
+(defn valid-totp?
+  "Returns true if the totp code is valid. Will go back up to max-10-minute-intervals."
+  ([^bytes secret-key max-5-minute-intervals ^String code]
+   (valid-totp? secret-key (or *now* (Instant/now)) max-5-minute-intervals code))
+  ([^bytes secret-key
+    ^Instant time
+    max-5-minute-intervals
+    ^String code]
+   (loop [remaining-intervals max-5-minute-intervals
+          time time]
+     (when (pos? remaining-intervals)
+       (if (crypt-util/constant-string= code (generate-totp secret-key time))
+         true
+         (recur (dec remaining-intervals)
+                (.minusSeconds time default-time-step)))))))

server/src/instant/util/crypt.clj

@@ -82,7 +82,7 @@
 (defn random-hex [^Long size]
   (bytes->hex-string (Random/randBytes size)))
 
-(defn hmac-256 [^bytes secret-key ^bytes b]
+(defn hmac-256 ^bytes [^bytes secret-key ^bytes b]
   (let [mac (Mac/getInstance "HmacSHA256")]
     (.init mac (SecretKeySpec. secret-key "HmacSHA256"))
     (.doFinal mac b)))