feat: External Migration

This commit is contained in:
Roi Feng
2025-06-18 18:20:07 -04:00
parent ae94e5a7ce
commit f32fb3836b
3 changed files with 53 additions and 2 deletions

View File

@ -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))
}

View File

@ -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 {

View File

@ -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))
}