diff --git a/internal/check/ExternalMigration.go b/internal/check/ExternalMigration.go index 76869e7..dcf4778 100644 --- a/internal/check/ExternalMigration.go +++ b/internal/check/ExternalMigration.go @@ -71,7 +71,7 @@ func genExternalMigrationSessionID(reqData dataType.UserRequest, ruleSet config. return []byte(fmt.Sprintf("%s:%s", fmt.Sprintf("%d", time.Now().Unix()), fmt.Sprintf("%x", mac.Sum(nil)))) } -func verifyExternalMigrationSessionIDCookie(reqData dataType.UserRequest, ruleSet config.RuleSet) bool { +func VerifyExternalMigrationSessionIDCookie(reqData dataType.UserRequest, ruleSet config.RuleSet) bool { if reqData.ToriiSessionID == "" { return false } @@ -95,3 +95,9 @@ func verifyExternalMigrationSessionIDCookie(reqData dataType.UserRequest, ruleSe return hmac.Equal([]byte(computedHash), []byte(expectedHash)) } + +func CalculateExternalMigrationHMAC(sessionID, timestampStr, secretKey string) string { + mac := hmac.New(sha512.New, []byte(secretKey)) + mac.Write([]byte(fmt.Sprintf("%s%s", sessionID, timestampStr))) + return fmt.Sprintf("%x", mac.Sum(nil)) +} diff --git a/internal/server/checker.go b/internal/server/checker.go index 6e6ba2f..2fc8a17 100644 --- a/internal/server/checker.go +++ b/internal/server/checker.go @@ -109,7 +109,7 @@ func CheckMain(w http.ResponseWriter, userRequestData dataType.UserRequest, rule } else if bytes.Compare(decision.HTTPCode, []byte("EXTERNAL")) == 0 { w.Header().Set("Set-Cookie", "__torii_sessionid="+string(decision.ResponseData)+"; Path=/; Max-Age=86400; Priority=High; HttpOnly; SameSite=Lax") - w.Header().Set("Location", ruleSet.ExternalMigrationRule.RedirectUrl+"?domain="+userRequestData.Host+"&session_id="+string(decision.ResponseData)+"&original_url="+userRequestData.Uri) + w.Header().Set("Location", ruleSet.ExternalMigrationRule.RedirectUrl+"?domain="+userRequestData.Host+"&session_id="+string(decision.ResponseData)+"&original_uri="+userRequestData.Uri) w.WriteHeader(http.StatusFound) return } else { diff --git a/internal/server/torii.go b/internal/server/torii.go index 830b8fa..e888b89 100644 --- a/internal/server/torii.go +++ b/internal/server/torii.go @@ -2,6 +2,7 @@ package server import ( "bytes" + "crypto/hmac" "html/template" "net/http" "server_torii/internal/action" @@ -9,6 +10,7 @@ import ( "server_torii/internal/config" "server_torii/internal/dataType" "server_torii/internal/utils" + "strconv" "time" ) @@ -20,6 +22,8 @@ func CheckTorii(w http.ResponseWriter, r *http.Request, reqData dataType.UserReq check.CheckCaptcha(r, reqData, ruleSet, decision) } else if reqData.Uri == cfg.WebPath+"/health_check" { decision.SetResponse(action.Done, []byte("200"), []byte("ok")) + } else if reqData.Uri == cfg.WebPath+"/external_migration" { + handleExternalMigration(w, r, reqData, ruleSet, decision) } if bytes.Compare(decision.HTTPCode, []byte("200")) == 0 { if bytes.Compare(decision.ResponseData, []byte("ok")) == 0 { @@ -88,3 +92,44 @@ func CheckTorii(w http.ResponseWriter, r *http.Request, reqData dataType.UserReq } } } + +func handleExternalMigration(w http.ResponseWriter, r *http.Request, data dataType.UserRequest, set *config.RuleSet, decision *action.Decision) { + if !set.ExternalMigrationRule.Enabled { + decision.SetResponse(action.Done, []byte("200"), []byte("bad")) + return + } + + originalURI := r.URL.Query().Get("original_uri") + timestampStr := r.URL.Query().Get("timestamp") + hmacParam := r.URL.Query().Get("hmac") + + if timestampStr == "" || hmacParam == "" { + decision.SetResponse(action.Done, []byte("200"), []byte("bad")) + return + } + + timestamp, err := strconv.ParseInt(timestampStr, 10, 64) + if err != nil { + decision.SetResponse(action.Done, []byte("200"), []byte("bad")) + return + } + + currentTime := time.Now().Unix() + if currentTime-timestamp > set.ExternalMigrationRule.SessionTimeout { + decision.SetResponse(action.Done, []byte("200"), []byte("bad")) + return + } + + if !check.VerifyExternalMigrationSessionIDCookie(data, *set) { + decision.SetResponse(action.Done, []byte("200"), []byte("badSession")) + return + } + + expectedHMAC := check.CalculateExternalMigrationHMAC(data.ToriiSessionID, timestampStr, set.ExternalMigrationRule.SecretKey) + if !hmac.Equal([]byte(expectedHMAC), []byte(hmacParam)) { + decision.SetResponse(action.Done, []byte("200"), []byte("bad")) + return + } + + decision.SetResponse(action.Continue, []byte("302"), []byte(originalURI)) +}