openclaw/tao-memory-mcp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: harden auth, sse, search, and docs

Browse Source
This commit is contained in:
OpenClaw Agent
2026-03-15 04:47:03 +08:00
parent f4574e6190
commit 80440077cf
4 changed files with 275 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

195
main.go
View File

@@ -8,6 +8,8 @@ import (
"log"
"net/http"
"os"
"path/filepath"
"strconv"
"strings"
"time"
)
@@ -90,14 +92,11 @@ func (s *TaoServer) dispatchMCP(token string, client string, req MCPRequest) {
return
}
if token == "" {
if token == "" && !getEnvBool("TAO_ALLOW_ANON", false) {
log.Printf("[MCP Response] missing token for method=%s", req.Method)
return
}
connKey := token
if client != "" {
connKey = token + "_" + client
}
connKey := buildConnKey(token, client)
if ch, ok := s.conns.Load(connKey); ok {
if b, err := json.Marshal(resp); err == nil {
ch.(chan string) <- string(b)
@@ -115,19 +114,120 @@ func getEnv(key, def string) string {
return def
}
func getEnvBool(key string, def bool) bool {
v := strings.ToLower(strings.TrimSpace(os.Getenv(key)))
if v == "" {
return def
}
switch v {
case "1", "true", "yes", "on":
return true
case "0", "false", "no", "off":
return false
default:
return def
}
}
func getEnvInt(key string, def int) int {
if v := os.Getenv(key); v != "" {
if n, err := strconv.Atoi(v); err == nil {
return n
}
}
return def
}
func extractToken(r *http.Request) (string, bool) {
if q := r.URL.Query().Get("token"); q != "" {
return q, true
}
h := r.Header.Get("Authorization")
if strings.HasPrefix(h, "Bearer ") {
return strings.TrimSpace(strings.TrimPrefix(h, "Bearer ")), false
}
return "", false
}
func buildConnKey(token string, client string) string {
if token == "" {
token = "anon"
}
if client != "" {
return token + "_" + client
}
return token
}
func generateClientID() string {
return fmt.Sprintf("c%d", time.Now().UnixNano())
}
func parseCORSOrigins() (bool, []string) {
raw := strings.TrimSpace(os.Getenv("TAO_CORS_ORIGINS"))
if raw == "" {
return false, nil
}
if raw == "*" {
return true, nil
}
parts := strings.Split(raw, ",")
var origins []string
for _, p := range parts {
p = strings.TrimSpace(p)
if p != "" {
origins = append(origins, p)
}
}
return false, origins
}
func setCORSHeaders(w http.ResponseWriter, r *http.Request) {
origin := r.Header.Get("Origin")
if origin == "" {
return
}
allowAll, origins := parseCORSOrigins()
if allowAll {
w.Header().Set("Access-Control-Allow-Origin", "*")
} else {
allowed := false
for _, o := range origins {
if o == origin {
allowed = true
break
}
}
if !allowed {
return
}
w.Header().Set("Access-Control-Allow-Origin", origin)
w.Header().Set("Vary", "Origin")
}
w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")
}
func isSubpath(path string, base string) bool {
absPath, err1 := filepath.Abs(path)
absBase, err2 := filepath.Abs(base)
if err1 != nil || err2 != nil {
return false
}
if absPath == absBase {
return true
}
return strings.HasPrefix(absPath, absBase+string(filepath.Separator))
}
// --- 以简御繁：鉴权 ---
func (s *TaoServer) checkAuth(r *http.Request) bool {
token := getEnv("TAO_AUTH_TOKEN", "")
if token == "" {
return true // 未配置则不启用鉴权
return getEnvBool("TAO_ALLOW_ANON", false)
}
// Header Bearer
h := r.Header.Get("Authorization")
if h == "Bearer "+token {
return true
}
// Query token
if q := r.URL.Query().Get("token"); q != "" && q == token {
reqToken, _ := extractToken(r)
if reqToken == token {
return true
}
return false
@@ -136,9 +236,7 @@ func (s *TaoServer) checkAuth(r *http.Request) bool {
func (s *TaoServer) requireAuth(next http.HandlerFunc) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
if r.Method == "OPTIONS" {
w.Header().Set("Access-Control-Allow-Origin", "*")
w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")
setCORSHeaders(w, r)
w.WriteHeader(http.StatusOK)
return
}
@@ -197,8 +295,8 @@ func (s *TaoServer) SSEHandler(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
w.Header().Set("Connection", "keep-alive")
w.Header().Set("Access-Control-Allow-Origin", "*")
w.Header().Set("X-Accel-Buffering", "no")
setCORSHeaders(w, r)
flusher, ok := w.(http.Flusher)
if !ok {
@@ -212,14 +310,24 @@ func (s *TaoServer) SSEHandler(w http.ResponseWriter, r *http.Request) {
if style == "message" {
endpoint = "message"
}
// 若通过 query token 访问，也把 token 拼到 endpoint（便于客户端无 Header）
token := r.URL.Query().Get("token")
queryToken := r.URL.Query().Get("token")
token, _ := extractToken(r)
if token == "" && !getEnvBool("TAO_ALLOW_ANON", false) {
http.Error(w, "Unauthorized", http.StatusUnauthorized)
return
}
client := r.URL.Query().Get("client")
if token != "" {
if client == "" {
client = generateClientID()
}
if queryToken != "" {
if strings.Contains(endpoint, "?") {
endpoint = endpoint + "&token=" + token
endpoint = endpoint + "&token=" + queryToken
} else {
endpoint = endpoint + "?token=" + token
endpoint = endpoint + "?token=" + queryToken
}
}
if client != "" {
@@ -232,16 +340,10 @@ func (s *TaoServer) SSEHandler(w http.ResponseWriter, r *http.Request) {
fmt.Fprintf(w, "event: endpoint\ndata: %s\n\n", endpoint)
flusher.Flush()
var msgChan chan string
if token != "" {
msgChan = make(chan string, 50)
connKey := token
if client != "" {
connKey = token + "_" + client
}
s.conns.Store(connKey, msgChan)
defer s.conns.Delete(connKey)
}
msgChan := make(chan string, 50)
connKey := buildConnKey(token, client)
s.conns.Store(connKey, msgChan)
defer s.conns.Delete(connKey)
ticker := time.NewTicker(5 * time.Second)
defer ticker.Stop()
@@ -262,9 +364,7 @@ func (s *TaoServer) SSEHandler(w http.ResponseWriter, r *http.Request) {
// --- MCP Message ---
func (s *TaoServer) MessageHandler(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Access-Control-Allow-Origin", "*")
w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")
setCORSHeaders(w, r)
if r.Method == "OPTIONS" {
w.WriteHeader(http.StatusOK)
@@ -273,7 +373,11 @@ func (s *TaoServer) MessageHandler(w http.ResponseWriter, r *http.Request) {
bodyBytes, _ := io.ReadAll(r.Body)
r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
log.Printf("[MCP POST] From=%s URL=%s Body=%s", r.RemoteAddr, r.URL.String(), string(bodyBytes))
if getEnvBool("TAO_DEBUG", false) {
log.Printf("[MCP POST] From=%s URL=%s Body=%s", r.RemoteAddr, r.URL.String(), string(bodyBytes))
} else {
log.Printf("[MCP POST] From=%s URL=%s", r.RemoteAddr, r.URL.String())
}
var req MCPRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
@@ -281,7 +385,7 @@ func (s *TaoServer) MessageHandler(w http.ResponseWriter, r *http.Request) {
return
}
token := r.URL.Query().Get("token")
token, _ := extractToken(r)
client := r.URL.Query().Get("client")
w.WriteHeader(http.StatusAccepted)
@@ -290,10 +394,23 @@ func (s *TaoServer) MessageHandler(w http.ResponseWriter, r *http.Request) {
// --- 主程序 (Main) ---
func main() {
if getEnv("TAO_AUTH_TOKEN", "") == "" && !getEnvBool("TAO_ALLOW_ANON", false) {
log.Fatal("TAO_AUTH_TOKEN is required unless TAO_ALLOW_ANON=true")
}
memoryRoot := getEnv("MEMORY_ROOT", "./knowledge_ocean")
searchRoot := getEnv("TAO_SEARCH_ROOT", memoryRoot)
if !isSubpath(searchRoot, memoryRoot) {
log.Printf("TAO_SEARCH_ROOT must be under MEMORY_ROOT, fallback to MEMORY_ROOT")
searchRoot = memoryRoot
}
server := &TaoServer{
config: Config{
MemoryRoot: getEnv("MEMORY_ROOT", "./knowledge_ocean"),
Port: getEnv("PORT", "5001"),
MemoryRoot: memoryRoot,
Port: getEnv("PORT", "5001"),
SearchRoot: searchRoot,
MaxSearchFiles: getEnvInt("TAO_SEARCH_MAX_FILES", 2000),
},
}