Allow sending resources on key challenge, send multiple challenges to specific browsers

lib/challenge.go

@@ -96,7 +96,7 @@ func (state *State) IssueChallengeToken(name string, key, result []byte, until t
 	return token, nil
 }
 
-func (state *State) VerifyChallengeToken(name string, expectedKey []byte, r *http.Request) (ok bool, err error) {
+func (state *State) VerifyChallengeToken(name string, expectedKey []byte, w http.ResponseWriter, r *http.Request) (ok bool, err error) {
 	c, ok := state.Challenges[name]
 	if !ok {
 		return false, errors.New("challenge not found")
@@ -104,7 +104,22 @@ func (state *State) VerifyChallengeToken(name string, expectedKey []byte, r *htt
 
 	cookie, err := r.Cookie(CookiePrefix + name)
 	if err != nil {
-		return false, err
+		// fallback: fetch cookie from response
+		if setCookies, ok := w.Header()["Set-Cookie"]; ok {
+			for _, setCookie := range setCookies {
+				newCookie, err := http.ParseSetCookie(setCookie)
+				if err != nil {
+					continue
+				}
+				// keep processing to find last set cookie
+				if newCookie.Name == name {
+					cookie = newCookie
+				}
+			}
+		}
+		if cookie == nil {
+			return false, err
+		}
 	}
 
 	token, err := jwt.ParseSigned(cookie.Value, []jose.SignatureAlgorithm{jose.EdDSA})

lib/http.go

@@ -143,7 +143,7 @@ func (state *State) handleRequest(w http.ResponseWriter, r *http.Request) {
 
 					for _, challengeName := range rule.Challenges {
 						key := state.GetChallengeKeyForRequest(challengeName, expiry, r)
-						ok, err := state.VerifyChallengeToken(challengeName, key, r)
+						ok, err := state.VerifyChallengeToken(challengeName, key, w, r)
 						if !ok || err != nil {
 							if !errors.Is(err, http.ErrNoCookie) {
 								ClearCookie(CookiePrefix+challengeName, w)

lib/state.go

@@ -347,6 +347,21 @@ func NewState(p policy.Policy, settings StateSettings) (state *State, err error)
 		case "":
 		case "http":
 		case "key":
+			mimeType := p.Parameters["key-mime"]
+			if mimeType == "" {
+				mimeType = "text/html; charset=utf-8"
+			}
+
+			httpCode, _ := strconv.Atoi(p.Parameters["key-code"])
+			if httpCode == 0 {
+				httpCode = http.StatusTemporaryRedirect
+			}
+
+			var content []byte
+			if data, ok := p.Parameters["key-content"]; ok {
+				content = []byte(data)
+			}
+
 			c.Verify = func(key []byte, result string) (bool, error) {
 				resultBytes, err := hex.DecodeString(result)
 				if err != nil {
@@ -359,6 +374,49 @@ func NewState(p policy.Policy, settings StateSettings) (state *State, err error)
 				return true, nil
 			}
 
+			c.VerifyChallenge = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+
+				err := func() (err error) {
+					expiry := time.Now().UTC().Add(DefaultValidity).Round(DefaultValidity)
+
+					key := state.GetChallengeKeyForRequest(challengeName, expiry, r)
+					result := r.FormValue("result")
+
+					if ok, err := c.Verify(key, result); err != nil {
+						return err
+					} else if !ok {
+						ClearCookie(CookiePrefix+challengeName, w)
+						_ = state.errorPage(w, http.StatusForbidden, fmt.Errorf("access denied: failed challenge %s", challengeName))
+						return nil
+					}
+
+					token, err := state.IssueChallengeToken(challengeName, key, []byte(result), expiry)
+					if err != nil {
+						ClearCookie(CookiePrefix+challengeName, w)
+					} else {
+						SetCookie(CookiePrefix+challengeName, token, expiry, w)
+					}
+
+					switch httpCode {
+					case http.StatusMovedPermanently, http.StatusFound, http.StatusSeeOther, http.StatusTemporaryRedirect, http.StatusPermanentRedirect:
+						http.Redirect(w, r, r.FormValue("redirect"), httpCode)
+					default:
+						w.Header().Set("Content-Type", mimeType)
+						w.WriteHeader(httpCode)
+						if content != nil {
+							_, _ = w.Write(content)
+						}
+					}
+
+					return nil
+				}()
+				if err != nil {
+					ClearCookie(CookiePrefix+challengeName, w)
+					_ = state.errorPage(w, http.StatusInternalServerError, err)
+					return
+				}
+			})
+
 		case "wasm":
 			wasmData, err := go_away.ChallengeFs.ReadFile(fmt.Sprintf("challenge/%s/runtime/%s", challengeName, p.Runtime.Asset))
 			if err != nil {