diff --git a/Makefile b/Makefile
index a1ae31188..4fca000e8 100755
--- a/Makefile
+++ b/Makefile
@@ -377,4 +377,4 @@ package-linux-arm-platform:
package-windows-platform:
@echo "Packaging Windows"
cd $(OUTPUT_DIR) && zip -r $(OUTPUT_DIR)/windows-amd64.zip $(APP_NAME)-windows-amd64.exe $(APP_CONFIG)
- cd $(OUTPUT_DIR) && zip -r $(OUTPUT_DIR)/windows-386.zip $(APP_NAME)-windows-386.exe $(APP_CONFIG)
\ No newline at end of file
+ cd $(OUTPUT_DIR) && zip -r $(OUTPUT_DIR)/windows-386.zip $(APP_NAME)-windows-386.exe $(APP_CONFIG)
diff --git a/cmd/vfs/main.go b/cmd/vfs/main.go
index dfe5f3593..d94456cbb 100755
--- a/cmd/vfs/main.go
+++ b/cmd/vfs/main.go
@@ -304,7 +304,7 @@ func (vfs StaticFS) Open(name string) (http.File, error) {
}
}
- log.Debug("local file not found,", localFile)
+ log.Trace("local file not found,", localFile)
}
if vfs.SkipVFS{
diff --git a/core/api/api.go b/core/api/api.go
index 4a0540ed4..c13674508 100755
--- a/core/api/api.go
+++ b/core/api/api.go
@@ -34,6 +34,7 @@ import (
"net"
"net/http"
"runtime"
+ "strings"
"sync"
"time"
@@ -120,7 +121,7 @@ func initializeAPI() {
}
// HandleAPIMethod register api handler
-func HandleAPIMethod(method Method, pattern string, handler func(w http.ResponseWriter, req *http.Request, ps httprouter.Params)) {
+func HandleAPIMethod(method Method, pattern string, handler func(w http.ResponseWriter, req *http.Request, ps httprouter.Params), options ...Option) {
l.Lock()
if registeredAPIMethodHandler == nil {
registeredAPIMethodHandler = map[string]map[string]func(w http.ResponseWriter, req *http.Request, ps httprouter.Params){}
@@ -132,8 +133,57 @@ func HandleAPIMethod(method Method, pattern string, handler func(w http.Response
registeredAPIMethodHandler[m] = map[string]func(w http.ResponseWriter, req *http.Request, ps httprouter.Params){}
}
registeredAPIMethodHandler[m][pattern] = handler
+ if len(options) > 0 {
+ opts := &HandlerOptions{}
+ for _, option := range options {
+ option(opts)
+ }
+ apiOptions.Register(method, pattern, opts)
+ }
+
+ l.Unlock()
+}
+func ServeRegisteredAPIRequest(w http.ResponseWriter, req *http.Request) {
+ localMux := http.NewServeMux()
+ localRouter := httprouter.New(localMux)
+ localRouter.NotFound = notfoundHandler
+
+ l.Lock()
+ funcHandlers := make(map[string]func(http.ResponseWriter, *http.Request), len(registeredAPIFuncHandler))
+ for pattern, handler := range registeredAPIFuncHandler {
+ funcHandlers[pattern] = handler
+ }
+ methodHandlers := make(map[string]map[string]func(w http.ResponseWriter, req *http.Request, ps httprouter.Params), len(registeredAPIMethodHandler))
+ for method, handlers := range registeredAPIMethodHandler {
+ cloned := make(map[string]func(w http.ResponseWriter, req *http.Request, ps httprouter.Params), len(handlers))
+ for pattern, handler := range handlers {
+ cloned[pattern] = handler
+ }
+ methodHandlers[method] = cloned
+ }
+ filterSnapshot := append([]filter.Filter(nil), filters...)
l.Unlock()
+
+ for pattern, handler := range funcHandlers {
+ wrapped := handler
+ for _, f := range filterSnapshot {
+ wrapped = f.FilterHttpHandlerFunc(pattern, wrapped)
+ }
+ localMux.HandleFunc(pattern, wrapped)
+ }
+
+ for method, handlers := range methodHandlers {
+ for pattern, handler := range handlers {
+ wrapped := handler
+ for _, f := range filterSnapshot {
+ wrapped = f.FilterHttpRouter(pattern, wrapped)
+ }
+ localRouter.Handle(method, pattern, wrapped)
+ }
+ }
+
+ localRouter.ServeHTTP(w, req)
}
var router = httprouter.New(mux)
@@ -145,8 +195,45 @@ var rootKey *rsa.PrivateKey
var rootCertPEM []byte
var apiConfig *config.APIConfig
-
var listenAddress string
+var resolveRuntimePublishIPv4 = util.GetIntranetIP
+
+func normalizeRuntimePublishAddress(actualAddr string) string {
+ actualAddr = strings.TrimSpace(actualAddr)
+ if actualAddr == "" {
+ return actualAddr
+ }
+
+ host, port, err := net.SplitHostPort(actualAddr)
+ if err != nil {
+ return actualAddr
+ }
+
+ normalizedHost := strings.Trim(strings.TrimSpace(host), "[]")
+ if normalizedHost != "" {
+ ip := net.ParseIP(normalizedHost)
+ if normalizedHost != util.AnyAddress && (ip == nil || !ip.IsUnspecified()) {
+ return actualAddr
+ }
+ }
+
+ ipv4, err := resolveRuntimePublishIPv4()
+ if err != nil || strings.TrimSpace(ipv4) == "" {
+ return actualAddr
+ }
+
+ return net.JoinHostPort(ipv4, port)
+}
+
+func syncRuntimePublishAddress(networkConfig *config.NetworkConfig, actualAddr string) {
+ if networkConfig == nil || strings.TrimSpace(actualAddr) == "" {
+ return
+ }
+ if strings.TrimSpace(networkConfig.Publish) != "" {
+ return
+ }
+ networkConfig.Publish = normalizeRuntimePublishAddress(actualAddr)
+}
var notfoundHandler = http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
rw.Write([]byte("{\"message\":\"not_found\"}"))
@@ -219,6 +306,7 @@ func StartAPI() {
if err != nil {
panic(err)
}
+ syncRuntimePublishAddress(&apiConfig.NetworkConfig, l.Addr().String())
router.NotFound = notfoundHandler
diff --git a/core/api/api_test.go b/core/api/api_test.go
index 45c7f3129..70ae597a9 100644
--- a/core/api/api_test.go
+++ b/core/api/api_test.go
@@ -27,9 +27,14 @@
package api
import (
+ "fmt"
"net/http"
"net/http/httptest"
"testing"
+ "time"
+
+ httprouter "infini.sh/framework/core/api/router"
+ "infini.sh/framework/core/config"
)
func TestStripPrefix(t *testing.T) {
@@ -111,3 +116,126 @@ func TestStripPrefix(t *testing.T) {
})
}
}
+
+func TestServeRegisteredAPIRequest(t *testing.T) {
+ path := fmt.Sprintf("/__copilot_test__/api/%s/:id", t.Name())
+ HandleAPIMethod(GET, path, func(w http.ResponseWriter, req *http.Request, ps httprouter.Params) {
+ w.WriteHeader(http.StatusAccepted)
+ _, _ = w.Write([]byte(ps.MustGetParameter("id") + ":" + req.URL.Query().Get("q")))
+ })
+
+ req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("%s/value?q=ok", fmt.Sprintf("/__copilot_test__/api/%s", t.Name())), nil)
+ recorder := httptest.NewRecorder()
+
+ ServeRegisteredAPIRequest(recorder, req)
+
+ if recorder.Code != http.StatusAccepted {
+ t.Fatalf("unexpected status: %d", recorder.Code)
+ }
+ if recorder.Body.String() != "value:ok" {
+ t.Fatalf("unexpected body: %s", recorder.Body.String())
+ }
+}
+
+func TestServeRegisteredAPIRequestAllowsNestedDispatch(t *testing.T) {
+ innerPath := fmt.Sprintf("/__copilot_test__/api/%s/inner", t.Name())
+ outerPath := fmt.Sprintf("/__copilot_test__/api/%s/outer", t.Name())
+
+ HandleAPIMethod(GET, innerPath, func(w http.ResponseWriter, req *http.Request, ps httprouter.Params) {
+ w.WriteHeader(http.StatusAccepted)
+ _, _ = w.Write([]byte("inner-ok"))
+ })
+ HandleAPIMethod(GET, outerPath, func(w http.ResponseWriter, req *http.Request, ps httprouter.Params) {
+ innerReq := httptest.NewRequest(http.MethodGet, innerPath, nil)
+ innerRecorder := httptest.NewRecorder()
+ ServeRegisteredAPIRequest(innerRecorder, innerReq)
+ w.WriteHeader(innerRecorder.Code)
+ _, _ = w.Write(innerRecorder.Body.Bytes())
+ })
+
+ req := httptest.NewRequest(http.MethodGet, outerPath, nil)
+ recorder := httptest.NewRecorder()
+
+ done := make(chan struct{})
+ go func() {
+ defer close(done)
+ ServeRegisteredAPIRequest(recorder, req)
+ }()
+
+ select {
+ case <-done:
+ case <-time.After(2 * time.Second):
+ t.Fatal("nested dispatch timed out")
+ }
+
+ if recorder.Code != http.StatusAccepted {
+ t.Fatalf("unexpected status: %d", recorder.Code)
+ }
+ if recorder.Body.String() != "inner-ok" {
+ t.Fatalf("unexpected body: %s", recorder.Body.String())
+ }
+}
+
+func TestSyncRuntimePublishAddressUsesActualListenAddressWhenUnset(t *testing.T) {
+ oldResolver := resolveRuntimePublishIPv4
+ resolveRuntimePublishIPv4 = func() (string, error) {
+ return "192.168.3.185", nil
+ }
+ t.Cleanup(func() {
+ resolveRuntimePublishIPv4 = oldResolver
+ })
+
+ cfg := config.NetworkConfig{}
+
+ syncRuntimePublishAddress(&cfg, "0.0.0.0:2901")
+
+ if cfg.Publish != "192.168.3.185:2901" {
+ t.Fatalf("expected runtime publish address to be updated, got %q", cfg.Publish)
+ }
+}
+
+func TestSyncRuntimePublishAddressNormalizesIPv6UnspecifiedHost(t *testing.T) {
+ oldResolver := resolveRuntimePublishIPv4
+ resolveRuntimePublishIPv4 = func() (string, error) {
+ return "192.168.3.185", nil
+ }
+ t.Cleanup(func() {
+ resolveRuntimePublishIPv4 = oldResolver
+ })
+
+ cfg := config.NetworkConfig{}
+
+ syncRuntimePublishAddress(&cfg, "[::]:2901")
+
+ if cfg.Publish != "192.168.3.185:2901" {
+ t.Fatalf("expected ipv6 unspecified runtime publish address to use ipv4, got %q", cfg.Publish)
+ }
+}
+
+func TestSyncRuntimePublishAddressPreservesExplicitPublishAddress(t *testing.T) {
+ cfg := config.NetworkConfig{Publish: "gateway.example:8443"}
+
+ syncRuntimePublishAddress(&cfg, "0.0.0.0:2901")
+
+ if cfg.Publish != "gateway.example:8443" {
+ t.Fatalf("expected explicit publish address to be preserved, got %q", cfg.Publish)
+ }
+}
+
+func TestSyncRuntimePublishAddressPreservesConcreteListenAddress(t *testing.T) {
+ oldResolver := resolveRuntimePublishIPv4
+ resolveRuntimePublishIPv4 = func() (string, error) {
+ return "192.168.3.185", nil
+ }
+ t.Cleanup(func() {
+ resolveRuntimePublishIPv4 = oldResolver
+ })
+
+ cfg := config.NetworkConfig{}
+
+ syncRuntimePublishAddress(&cfg, "10.0.0.8:2901")
+
+ if cfg.Publish != "10.0.0.8:2901" {
+ t.Fatalf("expected concrete runtime publish address to be preserved, got %q", cfg.Publish)
+ }
+}
diff --git a/core/api/basic_auth.go b/core/api/basic_auth.go
index 934720378..763bb49ea 100644
--- a/core/api/basic_auth.go
+++ b/core/api/basic_auth.go
@@ -28,8 +28,13 @@
package api
import (
+ "crypto/subtle"
httprouter "infini.sh/framework/core/api/router"
"net/http"
+ "strings"
+
+ "infini.sh/framework/core/model"
+ configcommon "infini.sh/framework/modules/configs/common"
)
type BasicAuthFilter struct {
@@ -37,9 +42,17 @@ type BasicAuthFilter struct {
Password string
}
+var loadManagedAccessTokenFromKeystore = func() (string, error) {
+ return configcommon.LoadTokenFromKeystore(configcommon.AgentAccessTokenKeystoreKey)
+}
+
// BasicAuth register api with basic auth
func BasicAuth(h httprouter.Handle, requiredUser, requiredPassword string) httprouter.Handle {
return func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ if validateManagedAccessToken(r) {
+ h(w, r, ps)
+ return
+ }
// Get the Basic Authentication credentials
user, password, hasAuth := r.BasicAuth()
@@ -60,6 +73,10 @@ func (filter *BasicAuthFilter) FilterHttpRouter(pattern string, h httprouter.Han
func (filter *BasicAuthFilter) FilterHttpHandlerFunc(pattern string, handler func(http.ResponseWriter, *http.Request)) func(http.ResponseWriter, *http.Request) {
return func(w http.ResponseWriter, request *http.Request) {
+ if validateManagedAccessToken(request) {
+ handler(w, request)
+ return
+ }
// Get the Basic Authentication credentials
user, password, hasAuth := request.BasicAuth()
if hasAuth && user == filter.Username && password == filter.Password {
@@ -72,3 +89,34 @@ func (filter *BasicAuthFilter) FilterHttpHandlerFunc(pattern string, handler fun
http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
}
}
+
+func validateManagedAccessToken(req *http.Request) bool {
+ tokenValue := ExtractBearerOrAPIToken(req)
+ if tokenValue == "" {
+ return false
+ }
+ expectedToken, err := loadManagedAccessTokenFromKeystore()
+ if err != nil || expectedToken == "" {
+ return false
+ }
+ return subtle.ConstantTimeCompare([]byte(expectedToken), []byte(tokenValue)) == 1
+}
+
+func ValidateManagedAccessTokenRequest(req *http.Request) bool {
+ return validateManagedAccessToken(req)
+}
+
+func ExtractBearerOrAPIToken(req *http.Request) string {
+ if req == nil {
+ return ""
+ }
+ tokenValue := strings.TrimSpace(req.Header.Get(model.API_TOKEN))
+ if tokenValue != "" {
+ return tokenValue
+ }
+ authHeader := strings.TrimSpace(req.Header.Get("Authorization"))
+ if len(authHeader) < len("Bearer ")+1 || !strings.EqualFold(authHeader[:len("Bearer ")], "Bearer ") {
+ return ""
+ }
+ return strings.TrimSpace(authHeader[len("Bearer "):])
+}
diff --git a/core/api/basic_auth_test.go b/core/api/basic_auth_test.go
new file mode 100644
index 000000000..d0912abbe
--- /dev/null
+++ b/core/api/basic_auth_test.go
@@ -0,0 +1,67 @@
+package api
+
+import (
+ "net/http"
+ "net/http/httptest"
+ "testing"
+
+ httprouter "infini.sh/framework/core/api/router"
+ "infini.sh/framework/core/model"
+)
+
+func TestBasicAuthAcceptsManagedAccessToken(t *testing.T) {
+ oldLoad := loadManagedAccessTokenFromKeystore
+ t.Cleanup(func() {
+ loadManagedAccessTokenFromKeystore = oldLoad
+ })
+ loadManagedAccessTokenFromKeystore = func() (string, error) {
+ return "managed-token", nil
+ }
+
+ handler := BasicAuth(func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ w.WriteHeader(http.StatusAccepted)
+ }, "api-user", "api-pass")
+
+ for name, applyAuth := range map[string]func(*http.Request){
+ "x-api-token": func(req *http.Request) {
+ req.Header.Set(model.API_TOKEN, "managed-token")
+ },
+ "bearer-token": func(req *http.Request) {
+ req.Header.Set("Authorization", "Bearer managed-token")
+ },
+ } {
+ t.Run(name, func(t *testing.T) {
+ req := httptest.NewRequest(http.MethodGet, "/stats", nil)
+ applyAuth(req)
+ recorder := httptest.NewRecorder()
+ handler(recorder, req, nil)
+
+ if recorder.Code != http.StatusAccepted {
+ t.Fatalf("unexpected status: %d", recorder.Code)
+ }
+ })
+ }
+}
+
+func TestBasicAuthFallsBackToBasicAuthCredentials(t *testing.T) {
+ oldLoad := loadManagedAccessTokenFromKeystore
+ t.Cleanup(func() {
+ loadManagedAccessTokenFromKeystore = oldLoad
+ })
+ loadManagedAccessTokenFromKeystore = func() (string, error) {
+ return "", nil
+ }
+
+ handler := BasicAuth(func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ w.WriteHeader(http.StatusAccepted)
+ }, "api-user", "api-pass")
+
+ req := httptest.NewRequest(http.MethodGet, "/stats", nil)
+ req.SetBasicAuth("api-user", "api-pass")
+ recorder := httptest.NewRecorder()
+ handler(recorder, req, nil)
+
+ if recorder.Code != http.StatusAccepted {
+ t.Fatalf("unexpected status: %d", recorder.Code)
+ }
+}
diff --git a/core/api/client.go b/core/api/client.go
index 6a9140d5f..c16bf177d 100755
--- a/core/api/client.go
+++ b/core/api/client.go
@@ -66,8 +66,10 @@ func SimpleGetTLSConfig(tlsConfig *config.TLSConfig) *tls.Config {
}
func GetClientTLSConfig(tlsConfig *config.TLSConfig) (*tls.Config, error) {
-
- pool := x509.NewCertPool()
+ pool, err := x509.SystemCertPool()
+ if err != nil || pool == nil {
+ pool = x509.NewCertPool()
+ }
skipVerify := tlsConfig.TLSInsecureSkipVerify
if tlsConfig.TLSBypassMalformedCert {
@@ -135,11 +137,10 @@ func GetClientTLSConfig(tlsConfig *config.TLSConfig) (*tls.Config, error) {
clientConfig.ServerName = "localhost"
}
- //skip domain verify if skip tls verify
- if !tlsConfig.TLSInsecureSkipVerify {
- if tlsConfig.SkipDomainVerify {
- clientConfig.VerifyPeerCertificate = util.GetSkipHostnameVerifyFunc(pool)
- }
+ // Skip hostname verification while still validating the certificate chain.
+ if tlsConfig.SkipDomainVerify && !tlsConfig.TLSInsecureSkipVerify {
+ clientConfig.InsecureSkipVerify = true
+ clientConfig.VerifyPeerCertificate = util.GetSkipHostnameVerifyFunc(pool)
}
return clientConfig, nil
diff --git a/core/api/client_test.go b/core/api/client_test.go
new file mode 100644
index 000000000..2f965e349
--- /dev/null
+++ b/core/api/client_test.go
@@ -0,0 +1,76 @@
+package api
+
+import (
+ "crypto/tls"
+ "net"
+ "os"
+ "path/filepath"
+ "testing"
+
+ "infini.sh/framework/core/config"
+ "infini.sh/framework/core/util"
+)
+
+func TestGetClientTLSConfigSkipDomainVerifyAllowsHostnameMismatch(t *testing.T) {
+ rootCert, rootKey, rootCertPEM := util.GetRootCert()
+ serverCertPEM, serverKeyPEM, err := util.GenerateServerCert(rootCert, rootKey, rootCertPEM, nil)
+ if err != nil {
+ t.Fatalf("generate server cert: %v", err)
+ }
+
+ dir := t.TempDir()
+ caFile := filepath.Join(dir, "ca.crt")
+ serverCertFile := filepath.Join(dir, "server.crt")
+ serverKeyFile := filepath.Join(dir, "server.key")
+
+ if err := os.WriteFile(caFile, rootCertPEM, 0600); err != nil {
+ t.Fatalf("write ca cert: %v", err)
+ }
+ if err := os.WriteFile(serverCertFile, serverCertPEM, 0600); err != nil {
+ t.Fatalf("write server cert: %v", err)
+ }
+ if err := os.WriteFile(serverKeyFile, serverKeyPEM, 0600); err != nil {
+ t.Fatalf("write server key: %v", err)
+ }
+
+ ln, err := net.Listen("tcp", "127.0.0.1:0")
+ if err != nil {
+ t.Fatalf("create listener: %v", err)
+ }
+ defer ln.Close()
+
+ serverCert, err := tls.LoadX509KeyPair(serverCertFile, serverKeyFile)
+ if err != nil {
+ t.Fatalf("load server cert: %v", err)
+ }
+ ln = tls.NewListener(ln, &tls.Config{Certificates: []tls.Certificate{serverCert}})
+
+ done := make(chan struct{})
+ go func() {
+ defer close(done)
+ conn, err := ln.Accept()
+ if err != nil {
+ return
+ }
+ if tlsConn, ok := conn.(*tls.Conn); ok {
+ _ = tlsConn.Handshake()
+ }
+ _ = conn.Close()
+ }()
+
+ cfg, err := GetClientTLSConfig(&config.TLSConfig{
+ TLSCACertFile: caFile,
+ SkipDomainVerify: true,
+ TLSInsecureSkipVerify: false,
+ })
+ if err != nil {
+ t.Fatalf("get client tls config: %v", err)
+ }
+
+ conn, err := tls.Dial("tcp", ln.Addr().String(), cfg)
+ if err != nil {
+ t.Fatalf("tls dial: %v", err)
+ }
+ _ = conn.Close()
+ <-done
+}
diff --git a/core/api/protected_routes.go b/core/api/protected_routes.go
new file mode 100644
index 000000000..01224ed2b
--- /dev/null
+++ b/core/api/protected_routes.go
@@ -0,0 +1,137 @@
+package api
+
+import (
+ "sort"
+
+ httprouter "infini.sh/framework/core/api/router"
+ "infini.sh/framework/core/util"
+)
+
+type ProtectedAPIRoute struct {
+ Method Method
+ Path string
+}
+
+var DefaultProtectedAPIRoutes = []ProtectedAPIRoute{
+ {Method: GET, Path: "/stats"},
+ {Method: GET, Path: "/queue/stats"},
+ {Method: GET, Path: "/queue/:id/stats"},
+ {Method: GET, Path: "/queue/:id/_scroll"},
+ {Method: DELETE, Path: "/queue/:id"},
+ {Method: DELETE, Path: "/queue/_search"},
+ {Method: PUT, Path: "/queue/:id/consumer/:consumer_id/offset"},
+ {Method: GET, Path: "/queue/:id/consumer/:consumer_id/offset"},
+ {Method: DELETE, Path: "/queue/:id/consumer/:consumer_id"},
+ {Method: DELETE, Path: "/queue/consumer/_search"},
+ {Method: GET, Path: "/pipeline/tasks/"},
+ {Method: POST, Path: "/pipeline/tasks/_search"},
+ {Method: POST, Path: "/pipeline/task/:id/_start"},
+ {Method: POST, Path: "/pipeline/task/:id/_stop"},
+ {Method: GET, Path: "/pipeline/task/:id"},
+ {Method: DELETE, Path: "/pipeline/task/:id"},
+ {Method: GET, Path: "/config/"},
+ {Method: PUT, Path: "/config/"},
+ {Method: GET, Path: "/config/runtime"},
+ {Method: GET, Path: "/setting/logger"},
+ {Method: PUT, Path: "/setting/logger"},
+ {Method: POST, Path: "/setting/logger"},
+}
+
+func RegisterProtectedUIRoutes(routes []ProtectedAPIRoute, handle httprouter.Handle, options ...Option) {
+ for _, route := range routes {
+ HandleUIMethod(route.Method, route.Path, handle, options...)
+ }
+}
+
+func RegisterProtectedRouterRoutes(router *httprouter.Router, routes []ProtectedAPIRoute, handle httprouter.Handle) {
+ if router == nil {
+ return
+ }
+ for _, route := range routes {
+ router.Handle(string(route.Method), route.Path, handle)
+ }
+}
+
+type MissingAPIMethodUIRoute struct {
+ Route ProtectedAPIRoute
+ Options *HandlerOptions
+}
+
+func WalkMissingAPIMethodUIRoutes(walk func(route MissingAPIMethodUIRoute)) {
+ if walk == nil {
+ return
+ }
+
+ l.Lock()
+ routes := make([]MissingAPIMethodUIRoute, 0)
+ for method, handlers := range registeredAPIMethodHandler {
+ for path := range handlers {
+ if shouldSkipEmbeddedAPIRoute(method, path) {
+ continue
+ }
+ var options *HandlerOptions
+ if registeredOptions, ok := apiOptions.Get(Method(method), path); ok {
+ options = cloneHandlerOptions(registeredOptions)
+ }
+ routes = append(routes, MissingAPIMethodUIRoute{
+ Route: ProtectedAPIRoute{
+ Method: Method(method),
+ Path: path,
+ },
+ Options: options,
+ })
+ }
+ }
+ l.Unlock()
+
+ sort.Slice(routes, func(i, j int) bool {
+ if routes[i].Route.Method == routes[j].Route.Method {
+ return routes[i].Route.Path < routes[j].Route.Path
+ }
+ return routes[i].Route.Method < routes[j].Route.Method
+ })
+
+ for _, route := range routes {
+ walk(route)
+ }
+}
+
+// RegisterMissingAPIMethodUIRoutes mirrors registered API method routes onto the
+// web router only when no UI route already owns the same method/path.
+func RegisterMissingAPIMethodUIRoutes(handle httprouter.Handle, options ...Option) {
+ if handle == nil {
+ return
+ }
+
+ WalkMissingAPIMethodUIRoutes(func(route MissingAPIMethodUIRoute) {
+ HandleUIMethod(route.Route.Method, route.Route.Path, handle, options...)
+ })
+}
+
+func cloneHandlerOptions(options *HandlerOptions) *HandlerOptions {
+ if options == nil {
+ return nil
+ }
+
+ cloned := *options
+ if options.RequirePermission != nil {
+ cloned.RequirePermission = append([]PermissionKey(nil), options.RequirePermission...)
+ }
+ if options.Tags != nil {
+ cloned.Tags = append([]string(nil), options.Tags...)
+ }
+ if options.Features != nil {
+ cloned.Features = map[string]bool{}
+ for key, value := range options.Features {
+ cloned.Features[key] = value
+ }
+ }
+ if options.Labels != nil {
+ cloned.Labels = util.MapStr{}
+ for key, value := range options.Labels {
+ cloned.Labels[key] = value
+ }
+ }
+
+ return &cloned
+}
diff --git a/core/api/security.go b/core/api/security.go
new file mode 100644
index 000000000..89a74cd48
--- /dev/null
+++ b/core/api/security.go
@@ -0,0 +1,155 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package api
+
+import (
+ "net/http"
+ "strings"
+
+ httprouter "infini.sh/framework/core/api/router"
+ replaysecurity "infini.sh/framework/core/security/replay"
+)
+
+type SecureTransportOptions struct {
+ // TrustForwardHeaders allows HTTPS detection to honor reverse-proxy forwarding headers.
+ TrustForwardHeaders bool
+}
+
+const (
+ // FeatureRequireSecureTransport marks a UI handler as HTTPS-only when it is enforced by filters.
+ FeatureRequireSecureTransport = "feature_require_secure_transport"
+ // FeatureRequireReplayProtection marks a UI handler as requiring a valid replay nonce.
+ FeatureRequireReplayProtection = "feature_require_replay_protection"
+ // LabelTrustForwardHeaders stores whether HTTPS checks may trust reverse-proxy forwarding headers.
+ LabelTrustForwardHeaders = "label_trust_forward_headers"
+)
+
+// RequestUsesSecureTransport reports whether the request arrived over HTTPS directly or, when
+// allowed, through a trusted reverse proxy that forwarded HTTPS metadata.
+func RequestUsesSecureTransport(req *http.Request, options ...SecureTransportOptions) bool {
+ if req == nil {
+ return false
+ }
+ if req.TLS != nil {
+ return true
+ }
+
+ resolved := resolveSecureTransportOptions(options)
+ if !resolved.TrustForwardHeaders {
+ return false
+ }
+
+ for _, header := range []string{"X-Forwarded-Proto", "X-Forwarded-Protocol", "X-Url-Scheme"} {
+ if headerIndicatesHTTPS(req.Header.Get(header)) {
+ return true
+ }
+ }
+
+ if strings.EqualFold(strings.TrimSpace(req.Header.Get("X-Forwarded-Ssl")), "on") {
+ return true
+ }
+
+ return forwardedHeaderIndicatesHTTPS(req.Header.Get("Forwarded"))
+}
+
+// RequireSecureTransport wraps a handler so it rejects requests that do not resolve to HTTPS.
+func (handler Handler) RequireSecureTransport(h httprouter.Handle, options ...SecureTransportOptions) httprouter.Handle {
+ resolved := resolveSecureTransportOptions(options)
+ return func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ if !RequestUsesSecureTransport(r, resolved) {
+ handler.WriteError(w, "this endpoint requires HTTPS. use https:// directly or route through a trusted HTTPS reverse proxy", http.StatusUpgradeRequired)
+ return
+ }
+ h(w, r, ps)
+ }
+}
+
+// RequireSecureTransport wraps a handler with the default security handler implementation.
+func RequireSecureTransport(h httprouter.Handle, options ...SecureTransportOptions) httprouter.Handle {
+ return Handler{}.RequireSecureTransport(h, options...)
+}
+
+// RequireReplayProtection wraps a handler so each request must present a valid replay nonce.
+func (handler Handler) RequireReplayProtection(h httprouter.Handle) httprouter.Handle {
+ return func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ if err := replaysecurity.ValidateAndConsumeReplayNonce(r); err != nil {
+ handler.WriteError(w, err.Error(), http.StatusUnauthorized)
+ return
+ }
+ h(w, r, ps)
+ }
+}
+
+// RequireReplayProtection wraps a handler with the default replay-protection implementation.
+func RequireReplayProtection(h httprouter.Handle) httprouter.Handle {
+ return Handler{}.RequireReplayProtection(h)
+}
+
+// SecureTransportOption annotates a UI route so SecurityFilter can enforce HTTPS consistently.
+func SecureTransportOption(options ...SecureTransportOptions) Option {
+ resolved := resolveSecureTransportOptions(options)
+ return func(o *HandlerOptions) {
+ Feature(FeatureRequireSecureTransport)(o)
+ Label(LabelTrustForwardHeaders, resolved.TrustForwardHeaders)(o)
+ }
+}
+
+// ReplayProtectionOption annotates a UI route so SecurityFilter enforces replay-nonce validation.
+func ReplayProtectionOption() Option {
+ return Feature(FeatureRequireReplayProtection)
+}
+
+func resolveSecureTransportOptions(options []SecureTransportOptions) SecureTransportOptions {
+ if len(options) == 0 {
+ return SecureTransportOptions{}
+ }
+ return options[0]
+}
+
+func headerIndicatesHTTPS(value string) bool {
+ if value == "" {
+ return false
+ }
+ first := strings.TrimSpace(strings.Split(value, ",")[0])
+ return strings.EqualFold(first, "https")
+}
+
+func forwardedHeaderIndicatesHTTPS(value string) bool {
+ if value == "" {
+ return false
+ }
+
+ for _, forwardedValue := range strings.Split(value, ",") {
+ for _, token := range strings.Split(forwardedValue, ";") {
+ parts := strings.SplitN(strings.TrimSpace(token), "=", 2)
+ if len(parts) != 2 || !strings.EqualFold(parts[0], "proto") {
+ continue
+ }
+ proto := strings.Trim(parts[1], "\"")
+ return strings.EqualFold(proto, "https")
+ }
+ }
+
+ return false
+}
diff --git a/core/api/security_test.go b/core/api/security_test.go
new file mode 100644
index 000000000..ce9002e09
--- /dev/null
+++ b/core/api/security_test.go
@@ -0,0 +1,160 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package api
+
+import (
+ "crypto/tls"
+ "net/http"
+ "net/http/httptest"
+ "testing"
+
+ httprouter "infini.sh/framework/core/api/router"
+ replaysecurity "infini.sh/framework/core/security/replay"
+)
+
+// The transport tests cover both direct TLS and trusted proxy headers because the
+// security helpers are shared by embedded UI routes that may sit behind a proxy.
+func TestRequestUsesSecureTransport(t *testing.T) {
+ tests := []struct {
+ name string
+ setup func(req *http.Request)
+ options []SecureTransportOptions
+ secure bool
+ }{
+ {
+ name: "tls request",
+ setup: func(req *http.Request) {
+ req.TLS = &tls.ConnectionState{}
+ },
+ secure: true,
+ },
+ {
+ name: "forwarded proto requires opt in",
+ setup: func(req *http.Request) {
+ req.Header.Set("X-Forwarded-Proto", "https")
+ },
+ secure: false,
+ },
+ {
+ name: "forwarded proto trusted when enabled",
+ setup: func(req *http.Request) {
+ req.Header.Set("X-Forwarded-Proto", "https")
+ },
+ options: []SecureTransportOptions{{TrustForwardHeaders: true}},
+ secure: true,
+ },
+ {
+ name: "plain http",
+ setup: func(req *http.Request) {},
+ secure: false,
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ req := httptest.NewRequest(http.MethodPost, "http://console.local/account/login", nil)
+ tt.setup(req)
+
+ if RequestUsesSecureTransport(req, tt.options...) != tt.secure {
+ t.Fatalf("expected secure=%v", tt.secure)
+ }
+ })
+ }
+}
+
+// The wrapper should fail fast before running the protected handler on plain HTTP.
+func TestRequireSecureTransport(t *testing.T) {
+ handler := Handler{}
+ called := false
+ protected := handler.RequireSecureTransport(func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ called = true
+ w.WriteHeader(http.StatusOK)
+ })
+
+ req := httptest.NewRequest(http.MethodPost, "http://console.local/account/login", nil)
+ resp := httptest.NewRecorder()
+
+ protected(resp, req, nil)
+
+ if called {
+ t.Fatal("expected insecure request to be blocked")
+ }
+ if resp.Code != http.StatusUpgradeRequired {
+ t.Fatalf("expected status %d, got %d", http.StatusUpgradeRequired, resp.Code)
+ }
+}
+
+// Replay-protected handlers should pass straight through once a matching nonce exists.
+func TestRequireReplayProtection(t *testing.T) {
+ handler := Handler{}
+ req := httptest.NewRequest(http.MethodPost, "https://console.local/account/login", nil)
+ nonce, _, err := replaysecurity.IssueReplayNonce(req, http.MethodPost, "/account/login")
+ if err != nil {
+ t.Fatalf("issue replay nonce: %v", err)
+ }
+ req.Header.Set(replaysecurity.HeaderName, nonce)
+
+ called := false
+ protected := handler.RequireReplayProtection(func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ called = true
+ w.WriteHeader(http.StatusOK)
+ })
+ resp := httptest.NewRecorder()
+
+ protected(resp, req, nil)
+
+ if !called {
+ t.Fatal("expected replay-protected handler to run")
+ }
+ if resp.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, resp.Code)
+ }
+}
+
+// Route options are later consumed by SecurityFilter, so the feature flag and labels
+// must both be set when secure transport enforcement is requested declaratively.
+func TestSecureTransportOption(t *testing.T) {
+ options := &HandlerOptions{}
+ SecureTransportOption(SecureTransportOptions{TrustForwardHeaders: true})(options)
+
+ if !options.Feature(FeatureRequireSecureTransport) {
+ t.Fatal("expected secure transport feature to be enabled")
+ }
+ if options.Labels == nil {
+ t.Fatal("expected labels to be initialized")
+ }
+ if v, ok := options.Labels[LabelTrustForwardHeaders].(bool); !ok || !v {
+ t.Fatalf("expected trust forward headers label to be true, got %#v", options.Labels[LabelTrustForwardHeaders])
+ }
+}
+
+// Replay protection uses a single feature flag because the filter reads no extra labels.
+func TestReplayProtectionOption(t *testing.T) {
+ options := &HandlerOptions{}
+ ReplayProtectionOption()(options)
+
+ if !options.Feature(FeatureRequireReplayProtection) {
+ t.Fatal("expected replay protection feature to be enabled")
+ }
+}
diff --git a/core/api/setting.go b/core/api/setting.go
index 5456ee24c..6b224aaf1 100644
--- a/core/api/setting.go
+++ b/core/api/setting.go
@@ -29,6 +29,7 @@ package api
import (
httprouter "infini.sh/framework/core/api/router"
+ "infini.sh/framework/core/global"
"infini.sh/framework/core/util"
"net/http"
"sync"
@@ -46,7 +47,8 @@ func init() {
func appSettingsAPIHandler(w http.ResponseWriter, req *http.Request, ps httprouter.Params) {
obj := util.MapStr{
- "auth_enabled": IsAuthEnable(),
+ "auth_enabled": IsAuthEnable(),
+ "setup_required": global.Env().SetupRequired(),
}
appSettings := GetAppSettings()
obj.Merge(appSettings)
diff --git a/core/api/setting_test.go b/core/api/setting_test.go
new file mode 100644
index 000000000..23b416be6
--- /dev/null
+++ b/core/api/setting_test.go
@@ -0,0 +1,38 @@
+package api
+
+import (
+ "encoding/json"
+ "net/http"
+ "net/http/httptest"
+ "testing"
+
+ "infini.sh/framework/core/env"
+ "infini.sh/framework/core/global"
+)
+
+func TestAppSettingsAPIHandlerIncludesSetupRequired(t *testing.T) {
+ oldEnv := global.Env()
+ testEnv := env.EmptyEnv()
+ testEnv.SystemConfig.PathConfig.Data = t.TempDir()
+ testEnv.EnableSetup(true)
+ global.RegisterEnv(testEnv)
+ defer global.RegisterEnv(oldEnv)
+
+ req := httptest.NewRequest(http.MethodGet, "/setting/application", nil)
+ resp := httptest.NewRecorder()
+
+ appSettingsAPIHandler(resp, req, nil)
+
+ if resp.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, resp.Code)
+ }
+
+ var body map[string]interface{}
+ if err := json.Unmarshal(resp.Body.Bytes(), &body); err != nil {
+ t.Fatalf("unmarshal response: %v", err)
+ }
+
+ if v, ok := body["setup_required"].(bool); !ok || !v {
+ t.Fatalf("expected setup_required=true, got %#v", body["setup_required"])
+ }
+}
diff --git a/core/api/web.go b/core/api/web.go
index 76406ff09..f8813f719 100755
--- a/core/api/web.go
+++ b/core/api/web.go
@@ -30,6 +30,8 @@ package api
import (
ctx "context"
"crypto/tls"
+ "errors"
+ "fmt"
"net/http"
_ "net/http/pprof"
"runtime"
@@ -56,16 +58,29 @@ var uiMutex sync.Mutex
var bindAddress string
+func ServeRegisteredUIRequest(w http.ResponseWriter, req *http.Request) error {
+ if uiRouter == nil {
+ return fmt.Errorf("web router is not initialized")
+ }
+ uiRouter.ServeHTTP(w, req)
+ return nil
+}
+
func StopWeb(cfg config.WebAppConfig) {
if srv != nil {
- ctx1, cancel := ctx.WithTimeout(ctx.Background(), 10*time.Second)
+ ctx1, cancel := ctx.WithTimeout(ctx.Background(), webShutdownTimeout)
defer cancel()
err := srv.Shutdown(ctx1)
if err != nil {
- panic(err)
+ log.Warnf("graceful web shutdown timed out or failed: %v, forcing close", err)
+ closeErr := srv.Close()
+ if closeErr != nil && !errors.Is(closeErr, http.ErrServerClosed) {
+ log.Errorf("force closing web server failed: %v", closeErr)
+ }
}
log.Debug("stopping web server")
+ srv = nil
}
}
@@ -110,6 +125,9 @@ func StartWeb(cfg config.WebAppConfig) {
if registeredAPIMethodHandler != nil {
for k, v := range registeredAPIMethodHandler {
for m, n := range v {
+ if shouldSkipEmbeddedAPIRoute(k, m) {
+ continue
+ }
log.Debug("register http handler: ", k, " ", m)
uiRouter.Handle(k, m, n)
}
@@ -117,6 +135,9 @@ func StartWeb(cfg config.WebAppConfig) {
}
if registeredAPIFuncHandler != nil {
for k, v := range registeredAPIFuncHandler {
+ if shouldSkipEmbeddedAPIRoute("", k) {
+ continue
+ }
log.Debug("register http handler: ", k)
uiServeMux.HandleFunc(k, v)
}
@@ -125,7 +146,10 @@ func StartWeb(cfg config.WebAppConfig) {
if cfg.WebsocketConfig.Enabled {
websocket.InitWebSocket(cfg.WebsocketConfig)
- uiServeMux.HandleFunc("/ws", websocket.ServeWs)
+ websocketPath := getWebsocketRegistrationPath(cfg)
+ if shouldRegisterWebsocketOnWeb(cfg) {
+ uiServeMux.HandleFunc(websocketPath, websocket.ServeWs)
+ }
if registeredWebSocketCommandHandler != nil {
for k, v := range registeredWebSocketCommandHandler {
log.Debug("register websocket handler: ", k, " ", v)
@@ -141,6 +165,7 @@ func StartWeb(cfg config.WebAppConfig) {
} else {
bindAddress = cfg.NetworkConfig.GetBindingAddr()
}
+ syncRuntimePublishAddress(&cfg.NetworkConfig, bindAddress)
handler := context.ClearHandler(uiRouter)
if cfg.Gzip.Enabled {
@@ -306,6 +331,48 @@ func (i *InterceptorHandler) AddInterceptors(interceptors ...Interceptor) {
}
}
+func getWebsocketRegistrationPath(cfg config.WebAppConfig) string {
+ if cfg.WebsocketConfig.BasePath != "" {
+ return cfg.WebsocketConfig.BasePath
+ }
+ return "/ws"
+}
+
+func shouldRegisterWebsocketOnWeb(cfg config.WebAppConfig) bool {
+ if !cfg.WebsocketConfig.Enabled {
+ return false
+ }
+ if !cfg.EmbeddingAPI || registeredAPIFuncHandler == nil {
+ return true
+ }
+ return registeredAPIFuncHandler[getWebsocketRegistrationPath(cfg)] == nil
+}
+
+func shouldSkipEmbeddedAPIRoute(method, path string) bool {
+ if registeredUIHandler != nil {
+ if _, exists := registeredUIHandler[path]; exists {
+ return true
+ }
+ }
+ if registeredUIMethodHandler == nil {
+ return false
+ }
+ if method == "" {
+ for _, handlers := range registeredUIMethodHandler {
+ if _, exists := handlers[path]; exists {
+ return true
+ }
+ }
+ return false
+ }
+ methodHandlers, exists := registeredUIMethodHandler[Method(method)]
+ if !exists {
+ return false
+ }
+ _, exists = methodHandlers[path]
+ return exists
+}
+
func (i *InterceptorHandler) Handler(handler http.Handler) http.Handler {
return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
var appliedInterceptors []Interceptor
@@ -346,6 +413,7 @@ func AddGlobalInterceptors(interceptors ...Interceptor) {
}
var srv *http.Server
+var webShutdownTimeout = 10 * time.Second
// RegisteredUIHandler is a hub for registered ui handler
var registeredUIHandler map[string]http.Handler
@@ -426,6 +494,7 @@ func HandleUIMethod(method Method, pattern string, handler func(w http.ResponseW
apiOptions.Register(method, pattern, opts)
}
+ _, hadPrevious := registeredUIMethodHandler[method][pattern]
if !opts.Override {
//check previous handler
previous, ok := registeredUIMethodHandler[method][pattern]
@@ -453,16 +522,28 @@ func HandleUIMethod(method Method, pattern string, handler func(w http.ResponseW
myHandler := RegisteredAPIHandler{Handler: handler, Options: opts}
registeredUIMethodHandler[method][pattern] = myHandler
+ registerLiveUIMethodHandler(method, pattern, myHandler, hadPrevious)
if opts.AllowOPTIONS {
m := registeredUIMethodHandler[OPTIONS]
+ hadOptionsPrevious := false
if m == nil {
registeredUIMethodHandler[OPTIONS] = map[string]RegisteredAPIHandler{}
+ } else {
+ _, hadOptionsPrevious = m[pattern]
}
registeredUIMethodHandler[OPTIONS][pattern] = myHandler
+ registerLiveUIMethodHandler(OPTIONS, pattern, myHandler, hadOptionsPrevious)
}
}
+func registerLiveUIMethodHandler(method Method, pattern string, handler RegisteredAPIHandler, alreadyRegistered bool) {
+ if uiRouter == nil || alreadyRegistered {
+ return
+ }
+ uiRouter.Handle(string(method), pattern, getWrappedHandler(string(method), pattern, handler))
+}
+
// HandleWebSocketCommand register websocket command handler
func HandleWebSocketCommand(command string, usage string, handler func(c *websocket.WebsocketConnection, array []string)) {
diff --git a/core/api/web_test.go b/core/api/web_test.go
new file mode 100644
index 000000000..399cd5fb3
--- /dev/null
+++ b/core/api/web_test.go
@@ -0,0 +1,258 @@
+package api
+
+import (
+ "net"
+ "net/http"
+ "net/http/httptest"
+ "testing"
+ "time"
+
+ httprouter "infini.sh/framework/core/api/router"
+ "infini.sh/framework/core/config"
+)
+
+func newTestBinding(t *testing.T) string {
+ t.Helper()
+
+ listener, err := net.Listen("tcp", "127.0.0.1:0")
+ if err != nil {
+ t.Fatalf("listen on random port: %v", err)
+ }
+ defer listener.Close()
+
+ return listener.Addr().String()
+}
+
+func TestWebsocketRegistrationPath(t *testing.T) {
+ cfg := config.WebAppConfig{}
+ cfg.WebsocketConfig.Enabled = true
+ cfg.WebsocketConfig.BasePath = "/custom-ws"
+
+ if got := getWebsocketRegistrationPath(cfg); got != "/custom-ws" {
+ t.Fatalf("unexpected websocket path: %s", got)
+ }
+
+ cfg.WebsocketConfig.BasePath = ""
+ if got := getWebsocketRegistrationPath(cfg); got != "/ws" {
+ t.Fatalf("unexpected default websocket path: %s", got)
+ }
+}
+
+func TestShouldRegisterWebsocketOnWeb(t *testing.T) {
+ originalHandlers := registeredAPIFuncHandler
+ t.Cleanup(func() {
+ registeredAPIFuncHandler = originalHandlers
+ })
+
+ cfg := config.WebAppConfig{}
+ cfg.WebsocketConfig.Enabled = true
+ cfg.WebsocketConfig.BasePath = "/ws"
+ cfg.EmbeddingAPI = true
+
+ registeredAPIFuncHandler = map[string]func(http.ResponseWriter, *http.Request){
+ "/ws": func(http.ResponseWriter, *http.Request) {},
+ }
+
+ if shouldRegisterWebsocketOnWeb(cfg) {
+ t.Fatal("expected embedded API websocket registration to suppress duplicate web registration")
+ }
+
+ delete(registeredAPIFuncHandler, "/ws")
+ if !shouldRegisterWebsocketOnWeb(cfg) {
+ t.Fatal("expected websocket registration when no embedded API websocket handler exists")
+ }
+
+ cfg.EmbeddingAPI = false
+ if !shouldRegisterWebsocketOnWeb(cfg) {
+ t.Fatal("expected websocket registration when embedding_api is disabled")
+ }
+}
+
+func TestShouldSkipEmbeddedAPIRoute(t *testing.T) {
+ originalUIHandlers := registeredUIHandler
+ originalUIMethodHandlers := registeredUIMethodHandler
+ t.Cleanup(func() {
+ registeredUIHandler = originalUIHandlers
+ registeredUIMethodHandler = originalUIMethodHandlers
+ })
+
+ registeredUIHandler = map[string]http.Handler{
+ "/": http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}),
+ }
+ registeredUIMethodHandler = map[Method]map[string]RegisteredAPIHandler{
+ GET: {
+ "/stats": {
+ Handler: func(http.ResponseWriter, *http.Request, httprouter.Params) {},
+ },
+ },
+ }
+
+ if !shouldSkipEmbeddedAPIRoute("", "/") {
+ t.Fatal("expected API root route to be skipped when UI root is registered")
+ }
+ if !shouldSkipEmbeddedAPIRoute(string(GET), "/stats") {
+ t.Fatal("expected method-based UI route to suppress embedded API registration")
+ }
+ if !shouldSkipEmbeddedAPIRoute("", "/stats") {
+ t.Fatal("expected UI method route to suppress embedded API func registration on same path")
+ }
+ if shouldSkipEmbeddedAPIRoute(string(GET), "/_info") {
+ t.Fatal("expected unrelated API route not to be skipped")
+ }
+}
+
+func TestRegisterMissingAPIMethodUIRoutesSkipsExistingUIRoutes(t *testing.T) {
+ originalAPIHandlers := registeredAPIMethodHandler
+ originalUIHandlers := registeredUIMethodHandler
+ originalServer := srv
+ originalRouter := uiRouter
+ originalServeMux := uiServeMux
+ t.Cleanup(func() {
+ registeredAPIMethodHandler = originalAPIHandlers
+ registeredUIMethodHandler = originalUIHandlers
+ srv = originalServer
+ uiRouter = originalRouter
+ uiServeMux = originalServeMux
+ })
+
+ registeredAPIMethodHandler = map[string]map[string]func(http.ResponseWriter, *http.Request, httprouter.Params){
+ http.MethodGet: {
+ "/api-only": func(w http.ResponseWriter, _ *http.Request, _ httprouter.Params) {
+ w.WriteHeader(http.StatusAccepted)
+ },
+ "/stats": func(w http.ResponseWriter, _ *http.Request, _ httprouter.Params) {
+ w.WriteHeader(http.StatusNoContent)
+ },
+ },
+ }
+ registeredUIMethodHandler = map[Method]map[string]RegisteredAPIHandler{
+ GET: {
+ "/stats": {
+ Handler: func(w http.ResponseWriter, _ *http.Request, _ httprouter.Params) {
+ w.WriteHeader(http.StatusCreated)
+ },
+ Options: &HandlerOptions{},
+ },
+ },
+ }
+
+ RegisterMissingAPIMethodUIRoutes(func(w http.ResponseWriter, _ *http.Request, _ httprouter.Params) {
+ w.WriteHeader(http.StatusOK)
+ })
+
+ webCfg := config.WebAppConfig{}
+ webCfg.NetworkConfig.Binding = newTestBinding(t)
+ StartWeb(webCfg)
+ defer StopWeb(webCfg)
+
+ resp := httptest.NewRecorder()
+ req := httptest.NewRequest(http.MethodGet, "/api-only", nil)
+ if err := ServeRegisteredUIRequest(resp, req); err != nil {
+ t.Fatalf("serve api-only ui route: %v", err)
+ }
+ if resp.Code != http.StatusOK {
+ t.Fatalf("expected missing API route to be mirrored onto web, got %d", resp.Code)
+ }
+
+ resp = httptest.NewRecorder()
+ req = httptest.NewRequest(http.MethodGet, "/stats", nil)
+ if err := ServeRegisteredUIRequest(resp, req); err != nil {
+ t.Fatalf("serve stats ui route: %v", err)
+ }
+ if resp.Code != http.StatusCreated {
+ t.Fatalf("expected existing UI route to win over mirrored API route, got %d", resp.Code)
+ }
+}
+
+func TestHandleUIMethodRegistersRouteAfterStartWeb(t *testing.T) {
+ originalUIHandlers := registeredUIMethodHandler
+ originalServer := srv
+ originalRouter := uiRouter
+ originalServeMux := uiServeMux
+ t.Cleanup(func() {
+ registeredUIMethodHandler = originalUIHandlers
+ srv = originalServer
+ uiRouter = originalRouter
+ uiServeMux = originalServeMux
+ })
+
+ registeredUIMethodHandler = map[Method]map[string]RegisteredAPIHandler{}
+
+ webCfg := config.WebAppConfig{}
+ webCfg.NetworkConfig.Binding = newTestBinding(t)
+ StartWeb(webCfg)
+ defer StopWeb(webCfg)
+
+ HandleUIMethod(GET, "/late-ui-route", func(w http.ResponseWriter, _ *http.Request, _ httprouter.Params) {
+ w.WriteHeader(http.StatusAccepted)
+ })
+
+ resp := httptest.NewRecorder()
+ req := httptest.NewRequest(http.MethodGet, "/late-ui-route", nil)
+ if err := ServeRegisteredUIRequest(resp, req); err != nil {
+ t.Fatalf("serve late ui route: %v", err)
+ }
+ if resp.Code != http.StatusAccepted {
+ t.Fatalf("expected late ui route to be available after web start, got %d", resp.Code)
+ }
+}
+
+func TestStopWebFallsBackToCloseWhenGracefulShutdownTimesOut(t *testing.T) {
+ originalUIHandlers := registeredUIMethodHandler
+ originalServer := srv
+ originalRouter := uiRouter
+ originalServeMux := uiServeMux
+ originalTimeout := webShutdownTimeout
+ t.Cleanup(func() {
+ registeredUIMethodHandler = originalUIHandlers
+ srv = originalServer
+ uiRouter = originalRouter
+ uiServeMux = originalServeMux
+ webShutdownTimeout = originalTimeout
+ })
+
+ registeredUIMethodHandler = map[Method]map[string]RegisteredAPIHandler{}
+ webCfg := config.WebAppConfig{}
+ webCfg.NetworkConfig.Binding = newTestBinding(t)
+ StartWeb(webCfg)
+
+ started := make(chan struct{})
+ release := make(chan struct{})
+ HandleUIMethod(GET, "/shutdown-timeout", func(w http.ResponseWriter, req *http.Request, _ httprouter.Params) {
+ close(started)
+ select {
+ case <-release:
+ case <-req.Context().Done():
+ }
+ w.WriteHeader(http.StatusOK)
+ })
+
+ clientDone := make(chan struct{})
+ go func() {
+ defer close(clientDone)
+ _, _ = http.Get("http://" + webCfg.NetworkConfig.Binding + "/shutdown-timeout")
+ }()
+
+ select {
+ case <-started:
+ case <-time.After(2 * time.Second):
+ t.Fatal("request did not reach blocking handler")
+ }
+
+ webShutdownTimeout = 50 * time.Millisecond
+
+ stopDone := make(chan struct{})
+ go func() {
+ defer close(stopDone)
+ StopWeb(webCfg)
+ }()
+
+ select {
+ case <-stopDone:
+ case <-time.After(2 * time.Second):
+ t.Fatal("StopWeb did not return after graceful shutdown timeout")
+ }
+
+ close(release)
+ <-clientDone
+}
diff --git a/core/api/websocket/conn.go b/core/api/websocket/conn.go
index 56e43c57f..77504d0c0 100755
--- a/core/api/websocket/conn.go
+++ b/core/api/websocket/conn.go
@@ -48,10 +48,12 @@ const (
// Send pings to peer with this period. Must be less than pongWait.
pingPeriod = (pongWait * 9) / 10
- // Maximum message size allowed from peer.
- maxMessageSize = 512
+ // Default maximum message size allowed from peer.
+ defaultMaxMessageSize int64 = 8 * 1024 * 1024
)
+var maxMessageSize int64 = defaultMaxMessageSize
+
var upgrader = websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
diff --git a/core/api/websocket/hub.go b/core/api/websocket/hub.go
index 068c3c473..386e5801c 100755
--- a/core/api/websocket/hub.go
+++ b/core/api/websocket/hub.go
@@ -87,6 +87,7 @@ func (h *Hub) registerHandlers() {
// InitWebSocket start websocket
func InitWebSocket(cfg config.WebsocketConfig) {
+ maxMessageSize = resolveMaxMessageSize(cfg)
if cfg.SkipHostVerify {
upgrader.CheckOrigin = func(r *http.Request) bool {
return true
@@ -123,6 +124,13 @@ func InitWebSocket(cfg config.WebsocketConfig) {
}
+func resolveMaxMessageSize(cfg config.WebsocketConfig) int64 {
+ if cfg.MaxMessageSizeBytes > 0 {
+ return cfg.MaxMessageSizeBytes
+ }
+ return defaultMaxMessageSize
+}
+
// HandleWebSocketCommand used to register command and handler
func HandleWebSocketCommand(cmd, usage string, handler func(c *WebsocketConnection, array []string)) {
cmd = strings.ToLower(strings.TrimSpace(cmd))
diff --git a/core/api/websocket/hub_test.go b/core/api/websocket/hub_test.go
new file mode 100644
index 000000000..380fd1faf
--- /dev/null
+++ b/core/api/websocket/hub_test.go
@@ -0,0 +1,36 @@
+package websocket
+
+import (
+ "testing"
+
+ "infini.sh/framework/core/config"
+)
+
+func TestResolveMaxMessageSize(t *testing.T) {
+ testCases := []struct {
+ name string
+ cfg config.WebsocketConfig
+ expect int64
+ }{
+ {
+ name: "default",
+ cfg: config.WebsocketConfig{},
+ expect: defaultMaxMessageSize,
+ },
+ {
+ name: "custom",
+ cfg: config.WebsocketConfig{
+ MaxMessageSizeBytes: 1024,
+ },
+ expect: 1024,
+ },
+ }
+
+ for _, tc := range testCases {
+ t.Run(tc.name, func(t *testing.T) {
+ if actual := resolveMaxMessageSize(tc.cfg); actual != tc.expect {
+ t.Fatalf("unexpected websocket message limit: got %d want %d", actual, tc.expect)
+ }
+ })
+ }
+}
diff --git a/core/api/websocket/reverse/manager.go b/core/api/websocket/reverse/manager.go
new file mode 100644
index 000000000..2a040e1e4
--- /dev/null
+++ b/core/api/websocket/reverse/manager.go
@@ -0,0 +1,312 @@
+package reverse
+
+import (
+ "bytes"
+ "context"
+ "errors"
+ "fmt"
+ "net/http"
+ "strings"
+ "sync"
+ "time"
+
+ "infini.sh/framework/core/util"
+)
+
+const (
+ DefaultTimeout = 30 * time.Second
+ DefaultMaxResponseBytes = 8 * 1024 * 1024
+ DefaultReconnectWait = 6 * time.Second
+ DefaultReconnectPoll = 200 * time.Millisecond
+)
+
+var (
+ ErrDisconnected = errors.New("reverse channel disconnected")
+ ErrNotConnected = errors.New("reverse channel is not connected")
+)
+
+type ManagerOptions struct {
+ DefaultTimeout time.Duration
+ MaxResponseBytes int
+ ReconnectWait time.Duration
+ ReconnectPoll time.Duration
+}
+
+type pendingResponse struct {
+ peerID string
+ body bytes.Buffer
+ status int
+ err error
+ done chan struct{}
+ completed bool
+}
+
+type SessionManager struct {
+ options ManagerOptions
+ mu sync.Mutex
+ pendingSessions map[string]string
+ activeSessions map[string]string
+ activeSessionsByID map[string]string
+ pendingResponses map[string]*pendingResponse
+}
+
+func NewSessionManager(options ManagerOptions) *SessionManager {
+ if options.DefaultTimeout <= 0 {
+ options.DefaultTimeout = DefaultTimeout
+ }
+ if options.MaxResponseBytes <= 0 {
+ options.MaxResponseBytes = DefaultMaxResponseBytes
+ }
+ if options.ReconnectWait <= 0 {
+ options.ReconnectWait = DefaultReconnectWait
+ }
+ if options.ReconnectPoll <= 0 {
+ options.ReconnectPoll = DefaultReconnectPoll
+ }
+ return &SessionManager{
+ options: options,
+ pendingSessions: map[string]string{},
+ activeSessions: map[string]string{},
+ activeSessionsByID: map[string]string{},
+ pendingResponses: map[string]*pendingResponse{},
+ }
+}
+
+func (m *SessionManager) RegisterPendingSession(sessionID, peerID string) {
+ m.mu.Lock()
+ defer m.mu.Unlock()
+ m.pendingSessions[sessionID] = strings.TrimSpace(peerID)
+}
+
+func (m *SessionManager) ActivateSession(sessionID, peerID string) error {
+ m.mu.Lock()
+ defer m.mu.Unlock()
+
+ peerID = strings.TrimSpace(peerID)
+ if expectedPeerID, ok := m.pendingSessions[sessionID]; !ok || expectedPeerID != peerID {
+ return fmt.Errorf("session handshake mismatch")
+ }
+ delete(m.pendingSessions, sessionID)
+
+ if previousSession, ok := m.activeSessions[peerID]; ok && previousSession != sessionID {
+ delete(m.activeSessionsByID, previousSession)
+ }
+
+ m.activeSessions[peerID] = sessionID
+ m.activeSessionsByID[sessionID] = peerID
+ return nil
+}
+
+func (m *SessionManager) HandleHelloPayload(payload string) error {
+ msg, err := ParseHelloPayload(payload)
+ if err != nil {
+ return err
+ }
+ return m.ActivateSession(msg.SessionID, msg.PeerID)
+}
+
+func (m *SessionManager) HandleResponsePayload(payload string) error {
+ msg, err := ParseResponsePayload(payload)
+ if err != nil {
+ return err
+ }
+ m.acceptResponse(msg)
+ return nil
+}
+
+func (m *SessionManager) OnDisconnect(sessionID string) {
+ m.mu.Lock()
+ defer m.mu.Unlock()
+
+ delete(m.pendingSessions, sessionID)
+ peerID, ok := m.activeSessionsByID[sessionID]
+ if !ok {
+ return
+ }
+
+ delete(m.activeSessionsByID, sessionID)
+ if currentSession, exists := m.activeSessions[peerID]; exists && currentSession == sessionID {
+ delete(m.activeSessions, peerID)
+ }
+ m.failPendingLocked(peerID, ErrDisconnected)
+}
+
+func (m *SessionManager) IsConnected(peerID string) bool {
+ m.mu.Lock()
+ defer m.mu.Unlock()
+ sessionID, ok := m.activeSessions[peerID]
+ return ok && sessionID != ""
+}
+
+func (m *SessionManager) WaitForReconnect(ctx context.Context, peerID string) bool {
+ waitCtx, cancel := context.WithTimeout(ctx, m.options.ReconnectWait)
+ defer cancel()
+
+ if m.IsConnected(peerID) {
+ return true
+ }
+
+ ticker := time.NewTicker(m.options.ReconnectPoll)
+ defer ticker.Stop()
+
+ for {
+ select {
+ case <-waitCtx.Done():
+ return false
+ case <-ticker.C:
+ if m.IsConnected(peerID) {
+ return true
+ }
+ }
+ }
+}
+
+func IsRecoverableError(err error) bool {
+ return errors.Is(err, ErrDisconnected) || errors.Is(err, ErrNotConnected)
+}
+
+func (m *SessionManager) ProxyRequest(peerID string, req *util.Request, headers http.Header, send func(sessionID, payload string) error, responseObjectToUnmarshal interface{}) (*util.Result, error) {
+ if req == nil {
+ return nil, fmt.Errorf("request is nil")
+ }
+
+ ctx := req.Context
+ if ctx == nil {
+ var cancel context.CancelFunc
+ ctx, cancel = context.WithTimeout(context.Background(), m.options.DefaultTimeout)
+ defer cancel()
+ } else if _, hasDeadline := ctx.Deadline(); !hasDeadline {
+ var cancel context.CancelFunc
+ ctx, cancel = context.WithTimeout(ctx, m.options.DefaultTimeout)
+ defer cancel()
+ }
+
+ var lastErr error
+ for attempt := 0; attempt < 2; attempt++ {
+ res, err := m.proxyRequestOnce(ctx, strings.TrimSpace(peerID), req, headers, send, responseObjectToUnmarshal)
+ if err == nil {
+ return res, nil
+ }
+ lastErr = err
+ if attempt == 0 && IsRecoverableError(err) && m.WaitForReconnect(ctx, peerID) {
+ continue
+ }
+ return res, err
+ }
+ return nil, lastErr
+}
+
+func (m *SessionManager) proxyRequestOnce(ctx context.Context, peerID string, req *util.Request, headers http.Header, send func(sessionID, payload string) error, responseObjectToUnmarshal interface{}) (*util.Result, error) {
+ requestID := util.GetUUID()
+ msg := RequestMessage{
+ RequestID: requestID,
+ PeerID: peerID,
+ Method: req.Method,
+ Path: req.Path,
+ Headers: headers,
+ }
+ msg.SetBody(req.Body)
+ if authorization := strings.TrimSpace(msg.Headers.Get("Authorization")); strings.HasPrefix(strings.ToLower(authorization), "bearer ") {
+ msg.AccessToken = strings.TrimSpace(authorization[7:])
+ }
+
+ pending := &pendingResponse{
+ peerID: peerID,
+ done: make(chan struct{}),
+ }
+
+ m.mu.Lock()
+ sessionID, ok := m.activeSessions[peerID]
+ if !ok || sessionID == "" {
+ m.mu.Unlock()
+ return nil, fmt.Errorf("%w for peer [%s]", ErrNotConnected, peerID)
+ }
+ m.pendingResponses[requestID] = pending
+ m.mu.Unlock()
+
+ if err := send(sessionID, FormatRequestCommand(msg)); err != nil {
+ m.mu.Lock()
+ delete(m.pendingResponses, requestID)
+ m.mu.Unlock()
+ return nil, err
+ }
+
+ select {
+ case <-pending.done:
+ case <-ctx.Done():
+ m.mu.Lock()
+ delete(m.pendingResponses, requestID)
+ m.mu.Unlock()
+ return nil, ctx.Err()
+ }
+
+ if pending.err != nil {
+ return nil, pending.err
+ }
+
+ res := &util.Result{
+ Body: pending.body.Bytes(),
+ StatusCode: pending.status,
+ }
+ if res.StatusCode != http.StatusOK {
+ return res, fmt.Errorf("request error: %s", string(res.Body))
+ }
+ if responseObjectToUnmarshal != nil && len(res.Body) > 0 {
+ return res, util.FromJSONBytes(res.Body, responseObjectToUnmarshal)
+ }
+ return res, nil
+}
+
+func (m *SessionManager) acceptResponse(msg ResponseMessage) {
+ m.mu.Lock()
+ defer m.mu.Unlock()
+
+ pending, ok := m.pendingResponses[msg.RequestID]
+ if !ok || pending.completed {
+ return
+ }
+ if msg.PeerID != "" && pending.peerID != "" && msg.PeerID != pending.peerID {
+ return
+ }
+
+ if msg.Chunk != "" {
+ chunk, err := msg.ChunkBytes()
+ if err != nil {
+ m.completePendingLocked(msg.RequestID, pending, 0, fmt.Errorf("decode reverse response chunk: %w", err))
+ return
+ }
+ if pending.body.Len()+len(chunk) > m.options.MaxResponseBytes {
+ m.completePendingLocked(msg.RequestID, pending, 0, fmt.Errorf("reverse response exceeds %d bytes", m.options.MaxResponseBytes))
+ return
+ }
+ _, _ = pending.body.Write(chunk)
+ }
+
+ if msg.Done {
+ status := msg.Status
+ if status == 0 {
+ status = http.StatusOK
+ }
+ m.completePendingLocked(msg.RequestID, pending, status, nil)
+ }
+}
+
+func (m *SessionManager) completePendingLocked(requestID string, pending *pendingResponse, status int, err error) {
+ if pending.completed {
+ return
+ }
+ pending.completed = true
+ pending.status = status
+ pending.err = err
+ close(pending.done)
+ delete(m.pendingResponses, requestID)
+}
+
+func (m *SessionManager) failPendingLocked(peerID string, err error) {
+ for requestID, pending := range m.pendingResponses {
+ if pending.peerID != peerID {
+ continue
+ }
+ m.completePendingLocked(requestID, pending, 0, err)
+ }
+}
diff --git a/core/api/websocket/reverse/manager_test.go b/core/api/websocket/reverse/manager_test.go
new file mode 100644
index 000000000..bc94cb97b
--- /dev/null
+++ b/core/api/websocket/reverse/manager_test.go
@@ -0,0 +1,73 @@
+package reverse
+
+import (
+ "net/http"
+ "strings"
+ "testing"
+
+ "infini.sh/framework/core/util"
+)
+
+func TestSessionManagerProxyRequestRoundTrip(t *testing.T) {
+ manager := NewSessionManager(ManagerOptions{})
+ manager.RegisterPendingSession("session-1", "peer-1")
+ if err := manager.ActivateSession("session-1", "peer-1"); err != nil {
+ t.Fatalf("activate session: %v", err)
+ }
+
+ headers := http.Header{}
+ headers.Set("Authorization", "Bearer token-1")
+
+ send := func(sessionID, payload string) error {
+ if sessionID != "session-1" {
+ t.Fatalf("unexpected session id: %s", sessionID)
+ }
+ if !strings.HasPrefix(payload, RequestCommand+" ") {
+ t.Fatalf("unexpected payload: %s", payload)
+ }
+ msg, err := ParseRequestPayload(strings.TrimPrefix(payload, RequestCommand+" "))
+ if err != nil {
+ t.Fatalf("parse request payload: %v", err)
+ }
+ if msg.BearerToken() != "token-1" {
+ t.Fatalf("unexpected bearer token: %s", msg.BearerToken())
+ }
+ return WriteChunkedResponse(func(responsePayload string) error {
+ if !strings.HasPrefix(responsePayload, ResponseCommand+" ") {
+ t.Fatalf("unexpected response payload: %s", responsePayload)
+ }
+ return manager.HandleResponsePayload(strings.TrimPrefix(responsePayload, ResponseCommand+" "))
+ }, msg.RequestID, msg.PeerID, http.StatusOK, []byte(`{"ack":true}`), DefaultResponseChunkBytes)
+ }
+
+ var response map[string]bool
+ req := &util.Request{Method: http.MethodGet, Path: "/stats"}
+ res, err := manager.ProxyRequest("peer-1", req, headers, send, &response)
+ if err != nil {
+ t.Fatalf("proxy request: %v", err)
+ }
+ if res.StatusCode != http.StatusOK {
+ t.Fatalf("unexpected status: %d", res.StatusCode)
+ }
+ if !response["ack"] {
+ t.Fatal("expected response to unmarshal")
+ }
+}
+
+func TestSessionManagerDisconnectFailsPendingRequest(t *testing.T) {
+ manager := NewSessionManager(ManagerOptions{})
+ manager.RegisterPendingSession("session-1", "peer-1")
+ if err := manager.ActivateSession("session-1", "peer-1"); err != nil {
+ t.Fatalf("activate session: %v", err)
+ }
+
+ send := func(sessionID, payload string) error {
+ manager.OnDisconnect(sessionID)
+ return nil
+ }
+
+ _, err := manager.ProxyRequest("peer-1", &util.Request{Method: http.MethodGet, Path: "/stats"}, nil, send, nil)
+ if !IsRecoverableError(err) {
+ t.Fatalf("expected recoverable disconnect error, got %v", err)
+ }
+}
diff --git a/core/api/websocket/reverse/protocol.go b/core/api/websocket/reverse/protocol.go
new file mode 100644
index 000000000..173337c88
--- /dev/null
+++ b/core/api/websocket/reverse/protocol.go
@@ -0,0 +1,162 @@
+package reverse
+
+import (
+ "encoding/base64"
+ "net/http"
+ "strings"
+
+ "infini.sh/framework/core/util"
+)
+
+const (
+ HeaderPeerID = "X-INFINI-INSTANCE-ID"
+ HelloCommand = "reverse_hello"
+ RequestCommand = "reverse_request"
+ ResponseCommand = "reverse_response"
+ DefaultResponseChunkBytes = 32 * 1024
+)
+
+type HelloMessage struct {
+ SessionID string `json:"session_id"`
+ PeerID string `json:"instance_id"`
+}
+
+type RequestMessage struct {
+ RequestID string `json:"request_id"`
+ PeerID string `json:"instance_id"`
+ Method string `json:"method"`
+ Path string `json:"path"`
+ Body string `json:"body,omitempty"`
+ Headers http.Header `json:"headers,omitempty"`
+ AccessToken string `json:"access_token,omitempty"`
+}
+
+type ResponseMessage struct {
+ RequestID string `json:"request_id"`
+ PeerID string `json:"instance_id"`
+ Chunk string `json:"chunk,omitempty"`
+ Status int `json:"status,omitempty"`
+ Done bool `json:"done,omitempty"`
+}
+
+func ParseHelloPayload(payload string) (HelloMessage, error) {
+ msg := HelloMessage{}
+ return msg, util.FromJSONBytes([]byte(payload), &msg)
+}
+
+func ParseRequestPayload(payload string) (RequestMessage, error) {
+ msg := RequestMessage{}
+ return msg, util.FromJSONBytes([]byte(payload), &msg)
+}
+
+func ParseResponsePayload(payload string) (ResponseMessage, error) {
+ msg := ResponseMessage{}
+ return msg, util.FromJSONBytes([]byte(payload), &msg)
+}
+
+func FormatHelloCommand(msg HelloMessage) string {
+ return HelloCommand + " " + string(util.MustToJSONBytes(msg))
+}
+
+func FormatRequestCommand(msg RequestMessage) string {
+ return RequestCommand + " " + string(util.MustToJSONBytes(msg))
+}
+
+func FormatResponseCommand(msg ResponseMessage) string {
+ return ResponseCommand + " " + string(util.MustToJSONBytes(msg))
+}
+
+func (m *RequestMessage) SetBody(body []byte) {
+ if len(body) == 0 {
+ m.Body = ""
+ return
+ }
+ m.Body = base64.StdEncoding.EncodeToString(body)
+}
+
+func (m RequestMessage) BodyBytes() ([]byte, error) {
+ if m.Body == "" {
+ return nil, nil
+ }
+ return base64.StdEncoding.DecodeString(m.Body)
+}
+
+func (m RequestMessage) NormalizedHeaders() http.Header {
+ headers := http.Header{}
+ for key, values := range m.Headers {
+ copied := append([]string(nil), values...)
+ headers[key] = copied
+ }
+ if headers.Get("Authorization") == "" && strings.TrimSpace(m.AccessToken) != "" {
+ headers.Set("Authorization", "Bearer "+strings.TrimSpace(m.AccessToken))
+ }
+ return headers
+}
+
+func (m RequestMessage) ApplyHeaders(req *http.Request) {
+ if req == nil {
+ return
+ }
+ if req.Header == nil {
+ req.Header = http.Header{}
+ }
+ for key := range req.Header {
+ req.Header.Del(key)
+ }
+ for key, values := range m.NormalizedHeaders() {
+ for _, value := range values {
+ req.Header.Add(key, value)
+ }
+ }
+}
+
+func (m RequestMessage) BearerToken() string {
+ value := strings.TrimSpace(m.NormalizedHeaders().Get("Authorization"))
+ if !strings.HasPrefix(strings.ToLower(value), "bearer ") {
+ return ""
+ }
+ return strings.TrimSpace(value[7:])
+}
+
+func (m *ResponseMessage) SetChunk(body []byte) {
+ if len(body) == 0 {
+ m.Chunk = ""
+ return
+ }
+ m.Chunk = base64.StdEncoding.EncodeToString(body)
+}
+
+func (m ResponseMessage) ChunkBytes() ([]byte, error) {
+ if m.Chunk == "" {
+ return nil, nil
+ }
+ return base64.StdEncoding.DecodeString(m.Chunk)
+}
+
+func WriteChunkedResponse(write func(payload string) error, requestID, peerID string, status int, body []byte, chunkBytes int) error {
+ if chunkBytes <= 0 {
+ chunkBytes = DefaultResponseChunkBytes
+ }
+ for start := 0; start < len(body); start += chunkBytes {
+ end := start + chunkBytes
+ if end > len(body) {
+ end = len(body)
+ }
+ msg := ResponseMessage{
+ RequestID: requestID,
+ PeerID: peerID,
+ }
+ msg.SetChunk(body[start:end])
+ if err := write(FormatResponseCommand(msg)); err != nil {
+ return err
+ }
+ }
+
+ done := ResponseMessage{
+ RequestID: requestID,
+ PeerID: peerID,
+ Status: status,
+ Done: true,
+ }
+ return write(FormatResponseCommand(done))
+}
diff --git a/core/api/websocket/reverse/protocol_test.go b/core/api/websocket/reverse/protocol_test.go
new file mode 100644
index 000000000..a2085d9da
--- /dev/null
+++ b/core/api/websocket/reverse/protocol_test.go
@@ -0,0 +1,43 @@
+package reverse
+
+import (
+ "net/http"
+ "testing"
+)
+
+func TestRequestMessageNormalizedHeadersFallsBackToLegacyAccessToken(t *testing.T) {
+ msg := RequestMessage{
+ AccessToken: "token-1",
+ }
+
+ headers := msg.NormalizedHeaders()
+ if got := headers.Get("Authorization"); got != "Bearer token-1" {
+ t.Fatalf("unexpected authorization header: %s", got)
+ }
+ if got := msg.BearerToken(); got != "token-1" {
+ t.Fatalf("unexpected bearer token: %s", got)
+ }
+}
+
+func TestRequestMessageApplyHeaders(t *testing.T) {
+ msg := RequestMessage{
+ Headers: http.Header{
+ "Authorization": []string{"Bearer token-2"},
+ "X-Test": []string{"value"},
+ },
+ }
+ req, _ := http.NewRequest(http.MethodGet, "http://example.com", nil)
+ req.Header.Set("Existing", "old")
+
+ msg.ApplyHeaders(req)
+
+ if req.Header.Get("Existing") != "" {
+ t.Fatal("expected old header to be removed")
+ }
+ if req.Header.Get("Authorization") != "Bearer token-2" {
+ t.Fatalf("unexpected authorization header: %s", req.Header.Get("Authorization"))
+ }
+ if req.Header.Get("X-Test") != "value" {
+ t.Fatalf("unexpected x-test header: %s", req.Header.Get("X-Test"))
+ }
+}
diff --git a/core/config/config.go b/core/config/config.go
index 2ddf8de3b..7d51b3ad1 100755
--- a/core/config/config.go
+++ b/core/config/config.go
@@ -270,7 +270,7 @@ func LoadEnvVariablesFromConfig(configObject *Config) (map[string]interface{}, e
return nil, err
}
- log.Debugf("config contain variables, try to parse with environments")
+ log.Tracef("config contains variables, parsing with environments")
environs := os.Environ()
obj := map[string]interface{}{}
@@ -345,7 +345,7 @@ func internalLoadFile(path string) (*Config, error) {
}
- log.Debugf("load config file '%v'", path)
+ log.Tracef("load config file '%v'", path)
return pCfg, err
}
diff --git a/core/config/fs_watcher.go b/core/config/fs_watcher.go
index 45706c402..873dd0700 100644
--- a/core/config/fs_watcher.go
+++ b/core/config/fs_watcher.go
@@ -62,6 +62,36 @@ func loadConfigFile(file string) *Config {
return nil
}
+func dispatchConfigChangeEvent(ev fsnotify.Event, watcherCallbacks []CallbackFunc) {
+ for _, v := range watcherCallbacks {
+ v(ev.Name, ev.Op)
+ }
+
+ cfg := loadConfigFile(ev.Name)
+ if cfg != nil {
+ for _, k := range sectionCallbackOrder {
+ callbacks, ok := sectionCallbacks[k]
+ if !ok || !cfg.HasField(k) {
+ continue
+ }
+ currentCfg, err := cfg.Child(k, -1)
+ if err != nil {
+ log.Error(err)
+ continue
+ }
+ previousCfg, _ := latestConfig[k]
+ for _, f := range callbacks {
+ f(previousCfg, currentCfg)
+ }
+ latestConfig[k] = currentCfg
+ }
+ }
+
+ for _, v := range configCallbacks {
+ v(ev)
+ }
+}
+
var validExtensions = []string{".yml", ".yaml", ".tpl"}
func SetValidExtension(v []string) {
@@ -153,40 +183,7 @@ func AddPathToWatch(path string, callback CallbackFunc) {
time.Sleep(2 * time.Second)
log.Trace("2 seconds out, on:", ev.String())
- // AddPathToWatch
-
- for _, v := range watcher.callbacks {
- v(ev.Name, ev.Op)
- }
-
- // NotifyOnConfigChange
-
- for _, v := range configCallbacks {
- v(ev)
- }
-
- // NotifyOnConfigSectionChange
-
- cfg := loadConfigFile(ev.Name)
- if cfg == nil {
- continue
- }
-
- for k, v := range sectionCallbacks {
- if cfg.HasField(k) {
- currentCfg, err := cfg.Child(k, -1)
- if err != nil {
- log.Error(err)
- continue
- }
- // diff config
- previousCfg, _ := latestConfig[k]
- for _, f := range v {
- f(previousCfg, currentCfg)
- }
- latestConfig[k] = currentCfg
- }
- }
+ dispatchConfigChangeEvent(ev, watcher.callbacks)
}
}()
})
@@ -255,11 +252,13 @@ func StopWatchers() {
}
var sectionCallbacks = map[string][]func(pCfg, cCfg *Config){}
+var sectionCallbackOrder = []string{}
var configCallbacks = []func(fsnotify.Event){}
var cfgLocker = sync.RWMutex{}
// NotifyOnConfigSectionChange will trigger callback when any configuration file change detected and
-// configKey present in the changed file
+// configKey present in the changed file. Section callbacks run before generic NotifyOnConfigChange
+// callbacks so section-scoped state can be refreshed before dependent consumers reload.
func NotifyOnConfigSectionChange(configKey string, f func(pCfg, cCfg *Config)) {
cfgLocker.Lock()
defer cfgLocker.Unlock()
@@ -268,12 +267,14 @@ func NotifyOnConfigSectionChange(configKey string, f func(pCfg, cCfg *Config)) {
if !ok {
v = []func(pCfg, cCfg *Config){}
sectionCallbacks[configKey] = v
+ sectionCallbackOrder = append(sectionCallbackOrder, configKey)
}
v = append(v, f)
sectionCallbacks[configKey] = v
}
-// NotifyOnConfigChange will trigger callback when any configuration file change detected
+// NotifyOnConfigChange will trigger callback when any configuration file change detected, after any
+// matching NotifyOnConfigSectionChange callbacks for the same event have run.
func NotifyOnConfigChange(f func(fsnotify.Event)) {
cfgLocker.Lock()
defer cfgLocker.Unlock()
diff --git a/core/config/fs_watcher_test.go b/core/config/fs_watcher_test.go
new file mode 100644
index 000000000..c1ddf1f75
--- /dev/null
+++ b/core/config/fs_watcher_test.go
@@ -0,0 +1,97 @@
+package config
+
+import (
+ "os"
+ "path/filepath"
+ "testing"
+
+ "github.com/fsnotify/fsnotify"
+)
+
+func TestDispatchConfigChangeEventRunsSectionCallbacksBeforeGenericCallbacks(t *testing.T) {
+ dir := t.TempDir()
+ file := filepath.Join(dir, "generated_metrics_tasks.yml")
+ content := []byte("elasticsearch:\n - id: \"cluster-1\"\n name: \"cluster-1\"\n enabled: true\n endpoint: \"http://127.0.0.1:9200\"\n")
+ if err := os.WriteFile(file, content, 0o644); err != nil {
+ t.Fatalf("write config file: %v", err)
+ }
+
+ previousSections := sectionCallbacks
+ previousOrder := sectionCallbackOrder
+ previousConfigs := configCallbacks
+ previousLatest := latestConfig
+ sectionCallbacks = map[string][]func(pCfg, cCfg *Config){}
+ sectionCallbackOrder = nil
+ configCallbacks = nil
+ latestConfig = map[string]*Config{}
+ t.Cleanup(func() {
+ sectionCallbacks = previousSections
+ sectionCallbackOrder = previousOrder
+ configCallbacks = previousConfigs
+ latestConfig = previousLatest
+ })
+
+ var order []string
+ NotifyOnConfigSectionChange("elasticsearch", func(pCfg, cCfg *Config) {
+ order = append(order, "section")
+ })
+ NotifyOnConfigChange(func(ev fsnotify.Event) {
+ order = append(order, "generic")
+ })
+
+ dispatchConfigChangeEvent(fsnotify.Event{Name: file, Op: fsnotify.Write}, nil)
+
+ if len(order) != 2 {
+ t.Fatalf("expected 2 callbacks, got %d (%v)", len(order), order)
+ }
+ if order[0] != "section" || order[1] != "generic" {
+ t.Fatalf("expected section callback before generic callback, got %v", order)
+ }
+}
+
+func TestDispatchConfigChangeEventRunsSectionCallbacksInRegistrationOrder(t *testing.T) {
+ dir := t.TempDir()
+ file := filepath.Join(dir, "gateway.yml")
+ content := []byte("flow:\n - name: flow-1\nrouter:\n - name: router-1\nentry:\n - name: entry-1\n")
+ if err := os.WriteFile(file, content, 0o644); err != nil {
+ t.Fatalf("write config file: %v", err)
+ }
+
+ previousSections := sectionCallbacks
+ previousOrder := sectionCallbackOrder
+ previousConfigs := configCallbacks
+ previousLatest := latestConfig
+ sectionCallbacks = map[string][]func(pCfg, cCfg *Config){}
+ sectionCallbackOrder = nil
+ configCallbacks = nil
+ latestConfig = map[string]*Config{}
+ t.Cleanup(func() {
+ sectionCallbacks = previousSections
+ sectionCallbackOrder = previousOrder
+ configCallbacks = previousConfigs
+ latestConfig = previousLatest
+ })
+
+ var order []string
+ NotifyOnConfigSectionChange("flow", func(pCfg, cCfg *Config) {
+ order = append(order, "flow")
+ })
+ NotifyOnConfigSectionChange("router", func(pCfg, cCfg *Config) {
+ order = append(order, "router")
+ })
+ NotifyOnConfigSectionChange("entry", func(pCfg, cCfg *Config) {
+ order = append(order, "entry")
+ })
+
+ dispatchConfigChangeEvent(fsnotify.Event{Name: file, Op: fsnotify.Write}, nil)
+
+ expected := []string{"flow", "router", "entry"}
+ if len(order) != len(expected) {
+ t.Fatalf("expected %d callbacks, got %d (%v)", len(expected), len(order), order)
+ }
+ for i, want := range expected {
+ if order[i] != want {
+ t.Fatalf("expected callback order %v, got %v", expected, order)
+ }
+ }
+}
diff --git a/core/config/system.go b/core/config/system.go
index b2f3bbc04..64fb8bee9 100755
--- a/core/config/system.go
+++ b/core/config/system.go
@@ -289,8 +289,9 @@ type ConfigsConfig struct {
ValidConfigsExtensions []string `config:"valid_config_extensions"`
TLSConfig TLSConfig `config:"tls"` //server or client's certs
ManagerConfig struct {
- LocalConfigsRepoPath string `config:"local_configs_repo_path"`
- BasicAuth BasicAuth `config:"basic_auth"`
+ LocalConfigsRepoPath string `config:"local_configs_repo_path"`
+ BasicAuth BasicAuth `config:"basic_auth"`
+ AccessToken ucfg.SecretString `config:"access_token"`
} `config:"manager"`
AlwaysRegisterAfterRestart bool `config:"always_register_after_restart"`
AllowGeneratedMetricsTasks bool `config:"allow_generated_metrics_tasks"`
@@ -379,6 +380,7 @@ type WebAppConfig struct {
//same with API Config
Enabled bool `config:"enabled"`
+ AccessLog bool `config:"access_log_enabled"`
TLSConfig TLSConfig `config:"tls"`
NetworkConfig NetworkConfig `config:"network"`
CrossDomain struct {
@@ -415,7 +417,7 @@ type S3BucketConfig struct {
}
func (config *WebAppConfig) GetEndpoint() string {
- return fmt.Sprintf("%s://%s", config.GetSchema(), config.NetworkConfig.GetPublishAddr())
+ return joinBasePath(fmt.Sprintf("%s://%s", config.GetSchema(), config.NetworkConfig.GetPublishAddr()), config.BasePath)
}
func (config *WebAppConfig) GetSchema() string {
@@ -451,7 +453,18 @@ type APIConfig struct {
}
func (config *APIConfig) GetEndpoint() string {
- return fmt.Sprintf("%s://%s", config.GetSchema(), config.NetworkConfig.GetPublishAddr())
+ return joinBasePath(fmt.Sprintf("%s://%s", config.GetSchema(), config.NetworkConfig.GetPublishAddr()), config.BasePath)
+}
+
+func joinBasePath(endpoint, basePath string) string {
+ basePath = strings.TrimSpace(basePath)
+ if basePath == "" || basePath == "/" {
+ return endpoint
+ }
+ if !strings.HasPrefix(basePath, "/") {
+ basePath = "/" + basePath
+ }
+ return strings.TrimRight(endpoint, "/") + strings.TrimRight(basePath, "/")
}
func (config *APIConfig) GetSchema() string {
@@ -511,6 +524,7 @@ type WebsocketConfig struct {
EchoWelcomeMessageOnConnect bool `config:"echo_welcome_message_on_connect"`
EchoLoggingConfigOnConnect bool `config:"echo_logging_config_on_connect"`
BasePath string `config:"base_path"`
+ MaxMessageSizeBytes int64 `config:"max_message_size_bytes"`
PermittedHosts []string `config:"permitted_hosts"`
SkipHostVerify bool `config:"skip_host_verify"`
}
diff --git a/core/config/system_test.go b/core/config/system_test.go
index 9224ec52d..7f2d24c5c 100644
--- a/core/config/system_test.go
+++ b/core/config/system_test.go
@@ -196,3 +196,25 @@ func TestHTTPClientConfig_ValidateProxy(t *testing.T) {
}
})
}
+
+func TestGetEndpointIncludesBasePath(t *testing.T) {
+ t.Run("api endpoint includes normalized base path", func(t *testing.T) {
+ cfg := APIConfig{
+ NetworkConfig: NetworkConfig{Publish: "agent.local:2900"},
+ BasePath: "api/v1/",
+ }
+ if got := cfg.GetEndpoint(); got != "http://agent.local:2900/api/v1" {
+ t.Fatalf("expected base path in api endpoint, got %q", got)
+ }
+ })
+
+ t.Run("web endpoint keeps root path unchanged", func(t *testing.T) {
+ cfg := WebAppConfig{
+ NetworkConfig: NetworkConfig{Publish: "console.local:9000"},
+ BasePath: "/",
+ }
+ if got := cfg.GetEndpoint(); got != "http://console.local:9000" {
+ t.Fatalf("expected root path to be ignored, got %q", got)
+ }
+ })
+}
diff --git a/core/credential/credential.go b/core/credential/credential.go
index 62e17a4e6..701ab99e1 100644
--- a/core/credential/credential.go
+++ b/core/credential/credential.go
@@ -31,6 +31,7 @@ import (
"fmt"
"infini.sh/framework/core/model"
"infini.sh/framework/core/orm"
+ "infini.sh/framework/lib/go-ucfg"
)
type Credential struct {
@@ -69,15 +70,20 @@ func (cred *Credential) Encode() error {
switch cred.Type {
case BasicAuth:
return encodeBasicAuth(cred)
+ case Token:
+ return encodeToken(cred)
+ case AccessToken:
+ return encodeAccessToken(cred)
default:
return fmt.Errorf("unkonow credential type [%s]", cred.Type)
}
}
+
func (cred *Credential) DecodeBasicAuth() (*model.BasicAuth, error) {
var dv interface{}
dv, err := cred.Decode()
if err != nil {
- panic(err)
+ return nil, err
}
if auth, ok := dv.(model.BasicAuth); ok {
@@ -86,15 +92,50 @@ func (cred *Credential) DecodeBasicAuth() (*model.BasicAuth, error) {
return nil, fmt.Errorf("unkonow credential type [%s]", cred.Type)
}
+func (cred *Credential) DecodeToken() (string, error) {
+ dv, err := cred.Decode()
+ if err != nil {
+ return "", err
+ }
+ if token, ok := dv.(model.Token); ok {
+ return token.Value, nil
+ }
+ return "", fmt.Errorf("unkonow credential type [%s]", cred.Type)
+}
+
+func (cred *Credential) DecodeAccessToken() (*AccessTokenPayload, error) {
+ dv, err := cred.Decode()
+ if err != nil {
+ return nil, err
+ }
+ if token, ok := dv.(AccessTokenPayload); ok {
+ return &token, nil
+ }
+ return nil, fmt.Errorf("unkonow credential type [%s]", cred.Type)
+}
+
func (cred *Credential) Decode() (interface{}, error) {
switch cred.Type {
case BasicAuth:
return decodeBasicAuth(cred)
+ case Token:
+ return decodeToken(cred)
+ case AccessToken:
+ return decodeAccessToken(cred)
default:
return nil, fmt.Errorf("unkonow credential type [%s]", cred.Type)
}
}
const (
- BasicAuth string = "basic_auth"
+ BasicAuth string = "basic_auth"
+ Token string = "token"
+ AccessToken string = "access_token"
)
+
+type AccessTokenPayload struct {
+ Value ucfg.SecretString `json:"value" yaml:"value"`
+ Permissions []string `json:"permissions,omitempty" yaml:"permissions,omitempty"`
+ Username string `json:"username,omitempty" yaml:"username,omitempty"`
+ Description string `json:"description,omitempty" yaml:"description,omitempty"`
+}
diff --git a/core/credential/domain.go b/core/credential/domain.go
index 2dbb05f4c..7459ef032 100644
--- a/core/credential/domain.go
+++ b/core/credential/domain.go
@@ -157,6 +157,142 @@ func decodeBasicAuth(cred *Credential) (basicAuth model.BasicAuth, err error) {
return
}
+func encodeToken(cred *Credential) error {
+ params, err := getCredentialPayloadMap(cred)
+ if err != nil {
+ return err
+ }
+ value, err := getCredentialSecret(params)
+ if err != nil {
+ return err
+ }
+ if value == "" {
+ return fmt.Errorf("credential parameters value can not be empty")
+ }
+
+ secret, err := GetOrInitSecret()
+ if err != nil {
+ return err
+ }
+ encodeBytes, salt, err := util.AesGcmEncrypt([]byte(value), secret)
+ if err != nil {
+ return fmt.Errorf("encrypt token value error: %w", err)
+ }
+ cred.Encrypt.Type = "AES"
+ cred.Encrypt.Params = map[string]interface{}{
+ "salt": string(salt),
+ }
+ params["value"] = string(encodeBytes)
+ cred.Payload[cred.Type] = params
+ return nil
+}
+
+func decodeToken(cred *Credential) (token model.Token, err error) {
+ params, err := getCredentialPayloadMap(cred)
+ if err != nil {
+ return
+ }
+ value, err := getCredentialSecret(params)
+ if err != nil {
+ return
+ }
+ if value == "" {
+ err = fmt.Errorf("credential parameters value can not be empty")
+ return
+ }
+ salt, ok := cred.Encrypt.Params["salt"].(string)
+ if !ok {
+ err = fmt.Errorf("credential encrypt parameters salt can not be empty")
+ return
+ }
+ secret := cred.secret
+ if secret == nil {
+ secret, err = GetOrInitSecret()
+ if err != nil {
+ return token, err
+ }
+ }
+ plaintext, err := util.AesGcmDecrypt([]byte(value), secret, []byte(salt))
+ if err != nil {
+ return token, err
+ }
+ token.Value = string(plaintext)
+ return
+}
+
+func encodeAccessToken(cred *Credential) error {
+ return encodeToken(cred)
+}
+
+func decodeAccessToken(cred *Credential) (payload AccessTokenPayload, err error) {
+ params, err := getCredentialPayloadMap(cred)
+ if err != nil {
+ return
+ }
+ value, err := getCredentialSecret(params)
+ if err != nil {
+ return
+ }
+ if value == "" {
+ err = fmt.Errorf("credential parameters value can not be empty")
+ return
+ }
+ salt, ok := cred.Encrypt.Params["salt"].(string)
+ if !ok {
+ err = fmt.Errorf("credential encrypt parameters salt can not be empty")
+ return
+ }
+ secret := cred.secret
+ if secret == nil {
+ secret, err = GetOrInitSecret()
+ if err != nil {
+ return payload, err
+ }
+ }
+ plaintext, err := util.AesGcmDecrypt([]byte(value), secret, []byte(salt))
+ if err != nil {
+ return payload, err
+ }
+ payload.Value = ucfg.SecretString(plaintext)
+ if permissions, ok := params["permissions"].([]interface{}); ok {
+ for _, item := range permissions {
+ if str, ok := item.(string); ok {
+ payload.Permissions = append(payload.Permissions, str)
+ }
+ }
+ } else if permissions, ok := params["permissions"].([]string); ok {
+ payload.Permissions = append(payload.Permissions, permissions...)
+ }
+ if username, ok := params["username"].(string); ok {
+ payload.Username = username
+ }
+ if description, ok := params["description"].(string); ok {
+ payload.Description = description
+ }
+ return
+}
+
+func getCredentialPayloadMap(cred *Credential) (map[string]interface{}, error) {
+ params, ok := cred.Payload[cred.Type].(map[string]interface{})
+ if !ok {
+ return nil, fmt.Errorf("wrong credential parameters for type [%s], expect a map", cred.Type)
+ }
+ return params, nil
+}
+
+func getCredentialSecret(params map[string]interface{}) (string, error) {
+ switch value := params["value"].(type) {
+ case string:
+ return value, nil
+ case []byte:
+ return string(value), nil
+ case ucfg.SecretString:
+ return string(value.Get()), nil
+ default:
+ return "", fmt.Errorf("wrong credential parameters value, expect a string")
+ }
+}
+
type ChangeEvent func(credentials *Credential)
var changeEvents []ChangeEvent
diff --git a/core/elastic/actions.go b/core/elastic/actions.go
index dcc11b03e..1eb11e5ef 100644
--- a/core/elastic/actions.go
+++ b/core/elastic/actions.go
@@ -118,10 +118,55 @@ func (node *NodeAvailable) IsDead() bool {
}
func (meta *ElasticsearchMetadata) IsAvailable() bool {
- if meta.Config == nil || !meta.Config.Enabled {
+ if meta.Config == nil {
+ if rate.GetRateLimiter("cluster_available_check", "nil_config", 1, 1, 30*time.Second).Allow() {
+ log.Debug("elasticsearch metadata is unavailable: config is nil")
+ }
+ return false
+ }
+ if !meta.Config.Enabled {
+ clusterID := meta.Config.ID
+ if clusterID == "" {
+ clusterID = meta.Config.Name
+ }
+ if rate.GetRateLimiter("cluster_available_check", clusterID, 1, 1, 30*time.Second).Allow() {
+ log.Debugf("elasticsearch [%v] is unavailable: config disabled", meta.Config.Name)
+ }
+ return false
+ }
+ if !meta.clusterAvailable {
+ clusterID := meta.Config.ID
+ if clusterID == "" {
+ clusterID = meta.Config.Name
+ }
+ if rate.GetRateLimiter("cluster_available_check", clusterID, 1, 1, 30*time.Second).Allow() {
+ if meta.shouldTraceUnavailableReason() {
+ log.Tracef("elasticsearch [%v] is unavailable: clusterAvailable=false", meta.Config.Name)
+ } else {
+ log.Debugf("elasticsearch [%v] is unavailable: clusterAvailable=false", meta.Config.Name)
+ }
+ }
return false
}
- return meta.clusterAvailable
+ return true
+}
+
+func (meta *ElasticsearchMetadata) shouldTraceUnavailableReason() bool {
+ if meta == nil || meta.Config == nil {
+ return false
+ }
+
+ return !meta.Config.Monitored
+}
+
+func (meta *ElasticsearchMetadata) shouldCheckActiveHostsOnFailure() bool {
+ if meta == nil || meta.Config == nil {
+ return true
+ }
+ if meta.Config.MetadataConfigs != nil && !meta.Config.MetadataConfigs.NodeAvailabilityCheck.Enabled {
+ return false
+ }
+ return true
}
func (meta *ElasticsearchMetadata) Init(health bool) {
@@ -186,13 +231,8 @@ func (meta *ElasticsearchMetadata) GetActiveEndpoint() string {
}
func (meta *ElasticsearchMetadata) GetActivePreferredSeedHost() string {
- hosts := meta.GetSeedHosts()
- if len(hosts) > 0 {
- for _, v := range hosts {
- if v != "" && IsHostAvailable(v) {
- return v
- }
- }
+ if host, _ := meta.getAvailableSeedHost(); host != "" {
+ return host
}
return meta.Config.Host
}
@@ -263,6 +303,12 @@ func (meta *ElasticsearchMetadata) GetActiveHosts() int {
}
func (meta *ElasticsearchMetadata) GetActiveHost() string {
+ if host, info := meta.getAvailableSeedHost(); host != "" {
+ if info != nil {
+ meta.activeHost = info
+ }
+ return host
+ }
if meta.activeHost != nil {
if meta.activeHost.IsAvailable() {
@@ -275,12 +321,9 @@ func (meta *ElasticsearchMetadata) GetActiveHost() string {
for _, v := range hosts {
if v != "" {
if IsHostAvailable(v) {
- //add to cache
- info, ok := GetHostAvailableInfo(v)
- if ok && info != nil {
- if info.IsAvailable() {
- meta.activeHost = info
- }
+ info := meta.ensureAvailableHostInfo(v)
+ if info != nil && info.IsAvailable() {
+ meta.activeHost = info
}
return v
@@ -295,12 +338,9 @@ func (meta *ElasticsearchMetadata) GetActiveHost() string {
v := v1.GetHttpPublishHost()
if v != "" {
if IsHostAvailable(v) {
- //add to cache
- info, ok := GetHostAvailableInfo(v)
- if ok && info != nil {
- if info.IsAvailable() {
- meta.activeHost = info
- }
+ info := meta.ensureAvailableHostInfo(v)
+ if info != nil && info.IsAvailable() {
+ meta.activeHost = info
}
return v
}
@@ -320,6 +360,46 @@ func (meta *ElasticsearchMetadata) GetActiveHost() string {
return hosts[0]
}
+func (meta *ElasticsearchMetadata) getAvailableSeedHost() (string, *NodeAvailable) {
+ hosts := meta.GetSeedHosts()
+ if hosts == nil || len(hosts) == 0 {
+ return "", nil
+ }
+
+ for _, host := range hosts {
+ if host == "" || !IsHostAvailable(host) {
+ continue
+ }
+ if info, ok := GetHostAvailableInfo(host); ok && info != nil && info.IsAvailable() {
+ return host, info
+ }
+ return host, meta.ensureAvailableHostInfo(host)
+ }
+
+ return "", nil
+}
+
+func (meta *ElasticsearchMetadata) ensureAvailableHostInfo(host string) *NodeAvailable {
+ if host == "" {
+ return nil
+ }
+
+ host = util.UnifyLocalAddress(host)
+ if info, ok := GetHostAvailableInfo(host); ok && info != nil {
+ return info
+ }
+
+ info := &NodeAvailable{
+ Host: host,
+ ClusterID: meta.Config.ID,
+ available: true,
+ lastCheck: time.Now(),
+ lastSuccess: time.Now(),
+ }
+ hosts.Store(host, info)
+ return info
+}
+
func (meta *ElasticsearchMetadata) IsTLS() bool {
return meta.GetSchema() == "https"
}
@@ -385,11 +465,15 @@ func (meta *ElasticsearchMetadata) ReportFailure(errorMessage error) bool {
return true
}
- num := meta.GetActiveHosts()
- log.Infof("%v has active hosts: %v", meta.Config.Name, num)
- if num > 0 {
- log.Debugf("enough failure ticket for elasticsearch [%v], but still have [%v] alive nodes", meta.Config.Name, num)
- return false
+ if meta.shouldCheckActiveHostsOnFailure() {
+ num := meta.GetActiveHosts()
+ log.Infof("%v has active hosts: %v", meta.Config.Name, num)
+ if num > 0 {
+ log.Debugf("enough failure ticket for elasticsearch [%v], but still have [%v] alive nodes", meta.Config.Name, num)
+ return false
+ }
+ } else if rate.GetRateLimiter("cluster_active_hosts_check", meta.Config.Name, 1, 1, 30*time.Second).Allow() {
+ log.Infof("skip active hosts check for elasticsearch [%v], node availability check is disabled", meta.Config.Name)
}
log.Debugf("enough failure ticket for elasticsearch [%v], mark it down", meta.Config.Name)
diff --git a/core/elastic/actions_test.go b/core/elastic/actions_test.go
new file mode 100644
index 000000000..5d706b190
--- /dev/null
+++ b/core/elastic/actions_test.go
@@ -0,0 +1,163 @@
+package elastic
+
+import (
+ "testing"
+ "time"
+
+ "infini.sh/framework/core/orm"
+)
+
+func TestGetActiveHostPrefersAvailableSeedHostOverCachedDiscoveredHost(t *testing.T) {
+ const (
+ clusterID = "docker-mapped-port-cluster"
+ seedHost = "192.168.3.185:9211"
+ discoveredHost = "172.18.1.18:9200"
+ )
+
+ cfg := &ElasticsearchConfig{
+ ORMObjectBase: orm.ORMObjectBase{ID: clusterID},
+ Name: clusterID,
+ Host: seedHost,
+ Hosts: []string{seedHost},
+ Enabled: true,
+ }
+ cfg.Discovery.Enabled = true
+
+ meta := &ElasticsearchMetadata{
+ Config: cfg,
+ Nodes: &map[string]NodesInfo{
+ "node-1": {
+ Http: struct {
+ BoundAddress []string `json:"bound_address"`
+ PublishAddress string `json:"publish_address,omitempty"`
+ MaxContentLengthInBytes int64 `json:"max_content_length_in_bytes,omitempty"`
+ }{
+ PublishAddress: discoveredHost,
+ },
+ },
+ },
+ activeHost: &NodeAvailable{Host: discoveredHost, available: true, lastCheck: time.Now()},
+ }
+
+ hosts.Store(seedHost, &NodeAvailable{Host: seedHost, ClusterID: clusterID, available: true, lastCheck: time.Now()})
+ hosts.Store(discoveredHost, &NodeAvailable{Host: discoveredHost, ClusterID: clusterID, available: true, lastCheck: time.Now()})
+ t.Cleanup(func() {
+ hosts.Delete(seedHost)
+ hosts.Delete(discoveredHost)
+ })
+
+ got := meta.GetActiveHost()
+ if got != seedHost {
+ t.Fatalf("expected seed host %q to be preferred over discovered host %q, got %q", seedHost, discoveredHost, got)
+ }
+ if meta.activeHost == nil || meta.activeHost.Host != seedHost {
+ t.Fatalf("expected activeHost to be updated to seed host %q, got %#v", seedHost, meta.activeHost)
+ }
+}
+
+func TestGetActiveHostFallsBackToCachedDiscoveredHostWhenSeedUnavailable(t *testing.T) {
+ const (
+ clusterID = "docker-discovery-fallback-cluster"
+ seedHost = "192.168.3.185:9211"
+ discoveredHost = "172.18.1.18:9200"
+ )
+
+ cfg := &ElasticsearchConfig{
+ ORMObjectBase: orm.ORMObjectBase{ID: clusterID},
+ Name: clusterID,
+ Host: seedHost,
+ Hosts: []string{seedHost},
+ Enabled: true,
+ }
+ cfg.Discovery.Enabled = true
+
+ meta := &ElasticsearchMetadata{
+ Config: cfg,
+ activeHost: &NodeAvailable{Host: discoveredHost, available: true, lastCheck: time.Now()},
+ }
+
+ hosts.Store(seedHost, &NodeAvailable{Host: seedHost, ClusterID: clusterID, available: false, lastCheck: time.Now()})
+ hosts.Store(discoveredHost, &NodeAvailable{Host: discoveredHost, ClusterID: clusterID, available: true, lastCheck: time.Now()})
+ t.Cleanup(func() {
+ hosts.Delete(seedHost)
+ hosts.Delete(discoveredHost)
+ })
+
+ got := meta.GetActiveHost()
+ if got != discoveredHost {
+ t.Fatalf("expected discovered host %q when seed host is unavailable, got %q", discoveredHost, got)
+ }
+}
+
+func TestGetActiveHostInitializesAvailableSeedHostInfoFromAvailabilityCache(t *testing.T) {
+ const (
+ clusterID = "seed-host-cache-init-cluster"
+ seedHost = "192.168.3.185:9220"
+ )
+
+ cfg := &ElasticsearchConfig{
+ ORMObjectBase: orm.ORMObjectBase{ID: clusterID},
+ Name: clusterID,
+ Host: seedHost,
+ Hosts: []string{seedHost},
+ Enabled: true,
+ }
+
+ meta := &ElasticsearchMetadata{Config: cfg}
+ nodeAvailCache.Put(seedHost, true)
+ hosts.Delete(seedHost)
+ t.Cleanup(func() {
+ hosts.Delete(seedHost)
+ })
+
+ got := meta.GetActiveHost()
+ if got != seedHost {
+ t.Fatalf("expected seed host %q, got %q", seedHost, got)
+ }
+
+ info, ok := GetHostAvailableInfo(seedHost)
+ if !ok || info == nil {
+ t.Fatalf("expected host info for %q to be initialized", seedHost)
+ }
+ if !info.IsAvailable() {
+ t.Fatalf("expected host info for %q to be marked available", seedHost)
+ }
+ if info.ClusterID != clusterID {
+ t.Fatalf("expected cluster id %q, got %q", clusterID, info.ClusterID)
+ }
+}
+
+func TestShouldTraceUnavailableReasonForUnmonitoredCluster(t *testing.T) {
+ meta := &ElasticsearchMetadata{
+ Config: &ElasticsearchConfig{
+ Monitored: false,
+ },
+ }
+
+ if !meta.shouldTraceUnavailableReason() {
+ t.Fatal("expected unmonitored cluster to trace unavailable reason")
+ }
+
+ meta.Config.Monitored = true
+ if meta.shouldTraceUnavailableReason() {
+ t.Fatal("expected monitored cluster to keep debug unavailable reason")
+ }
+}
+
+func TestShouldCheckActiveHostsOnFailure(t *testing.T) {
+ meta := &ElasticsearchMetadata{Config: &ElasticsearchConfig{}}
+ if !meta.shouldCheckActiveHostsOnFailure() {
+ t.Fatal("expected active hosts check enabled by default")
+ }
+
+ meta.Config.MetadataConfigs = &MetadataConfig{}
+ meta.Config.MetadataConfigs.NodeAvailabilityCheck.Enabled = true
+ if !meta.shouldCheckActiveHostsOnFailure() {
+ t.Fatal("expected active hosts check enabled when node availability check is on")
+ }
+
+ meta.Config.MetadataConfigs.NodeAvailabilityCheck.Enabled = false
+ if meta.shouldCheckActiveHostsOnFailure() {
+ t.Fatal("expected active hosts check disabled when node availability check is off")
+ }
+}
diff --git a/core/elastic/common_command.go b/core/elastic/common_command.go
index bc2cbdccf..0a4e9cd91 100644
--- a/core/elastic/common_command.go
+++ b/core/elastic/common_command.go
@@ -35,6 +35,7 @@ type CommonCommand struct {
ID string `json:"-" index:"id"`
Title string `json:"title" elastic_mapping:"title:{type:text,fields:{keyword:{type:keyword}}}"`
Tag []string `json:"tag" elastic_mapping:"tag:{type:keyword}"`
+ Creator string `json:"creator,omitempty" elastic_mapping:"creator:{type:keyword}"`
Requests []CommandRequest `json:"requests" elastic_mapping:"requests:{type:object}"`
Created time.Time `json:"created,omitempty" elastic_mapping:"created:{type:date}"`
}
diff --git a/core/elastic/domain.go b/core/elastic/domain.go
index 97bb08ca0..bcac81c46 100644
--- a/core/elastic/domain.go
+++ b/core/elastic/domain.go
@@ -541,6 +541,9 @@ type ElasticsearchConfig struct {
Distribution string `json:"distribution,omitempty" elastic_mapping:"distribution:{type:keyword}"`
NoDefaultAuthForAgent bool `json:"no_default_auth_for_agent,omitempty" config:"no_default_auth_for_agent"`
MetricCollectionMode string `json:"metric_collection_mode,omitempty" elastic_mapping:"metric_collection_mode:{type:keyword}"`
+ // AgentCollectionInterval is the default metrics collection interval (in seconds) for all Agent pipelines
+ // monitoring this cluster. 0 means use the Agent binary default (10 s). Can be overridden per-node via node_settings.
+ AgentCollectionInterval int `json:"agent_collection_interval,omitempty" elastic_mapping:"agent_collection_interval:{type:integer}"`
}
const (
diff --git a/core/elastic/domain_actions.go b/core/elastic/domain_actions.go
index d40bceadf..634246a08 100644
--- a/core/elastic/domain_actions.go
+++ b/core/elastic/domain_actions.go
@@ -43,6 +43,7 @@ import (
"crypto/tls"
"fmt"
uri "net/url"
+ "strconv"
"strings"
"sync"
"time"
@@ -99,8 +100,14 @@ func RegisterInstance(cfg ElasticsearchConfig, handler API) {
UpdateClient(cfg, handler)
UpdateConfig(cfg)
+ meta := GetMetadata(cfg.ID)
+ if meta == nil {
+ InitMetadata(&cfg, false)
+ return
+ }
+
if exists && oldCfg != nil {
- InitMetadata(&cfg, true)
+ InitMetadata(&cfg, meta.IsAvailable())
}
}
@@ -202,6 +209,54 @@ func (c *ElasticsearchConfig) GetAnyEndpoint() string {
panic(fmt.Errorf("no endpoint was not found in config [%v] ", c.ID))
}
+func (c *ElasticsearchConfig) GetAllEndpoints() []string {
+ build := func(host string) string {
+ return fmt.Sprintf("%s://%s", c.Schema, host)
+ }
+
+ seen := make(map[string]struct{})
+ result := make([]string, 0)
+
+ quote := func(v string) string {
+ if v == "" {
+ return ""
+ }
+ return strconv.Quote(v)
+ }
+
+ add := func(v string) {
+ if v == "" {
+ return
+ }
+
+ qv := quote(v)
+
+ if _, ok := seen[qv]; ok {
+ return
+ }
+ seen[qv] = struct{}{}
+ result = append(result, qv)
+ }
+
+ // 1. Hosts -> schema + host
+ for _, host := range c.Hosts {
+ add(build(host))
+ }
+
+ // 2. Endpoints -> raw
+ for _, ep := range c.Endpoints {
+ add(ep)
+ }
+
+ // 3. Endpoint -> raw single
+ add(c.Endpoint)
+
+ // 4. Host -> schema + host
+ add(build(c.Host))
+
+ return result
+}
+
func (meta *ElasticsearchMetadata) GetMajorVersion() int {
versionLock.RLock()
diff --git a/core/elastic/domain_actions_test.go b/core/elastic/domain_actions_test.go
new file mode 100644
index 000000000..1a2c05a56
--- /dev/null
+++ b/core/elastic/domain_actions_test.go
@@ -0,0 +1,35 @@
+package elastic
+
+import (
+ "testing"
+
+ "infini.sh/framework/core/orm"
+)
+
+func TestRegisterInstanceInitializesMetadataOnFirstRegistration(t *testing.T) {
+ cfg := ElasticsearchConfig{
+ ORMObjectBase: orm.ORMObjectBase{ID: "test-first-sync"},
+ Name: "test-first-sync",
+ Enabled: true,
+ ClusterUUID: "cluster-uuid-1",
+ }
+
+ t.Cleanup(func() {
+ cfgs.Delete(cfg.ID)
+ apis.Delete(cfg.ID)
+ metas.Delete(cfg.ID)
+ })
+
+ RegisterInstance(cfg, nil)
+
+ meta := GetMetadata(cfg.ID)
+ if meta == nil {
+ t.Fatalf("expected metadata to be initialized for %s", cfg.ID)
+ }
+ if meta.Config == nil {
+ t.Fatalf("expected metadata config to be initialized for %s", cfg.ID)
+ }
+ if meta.Config.ClusterUUID != cfg.ClusterUUID {
+ t.Fatalf("expected cluster uuid %q, got %q", cfg.ClusterUUID, meta.Config.ClusterUUID)
+ }
+}
diff --git a/core/elastic/index.go b/core/elastic/index.go
index 6ca476b22..bdb034440 100755
--- a/core/elastic/index.go
+++ b/core/elastic/index.go
@@ -24,10 +24,13 @@
package elastic
import (
+ "bytes"
"errors"
"github.com/buger/jsonparser"
"github.com/segmentio/encoding/json"
"infini.sh/framework/core/util"
+ "sort"
+ "strconv"
"strings"
"time"
)
@@ -213,8 +216,72 @@ type Bucket struct {
}
type AggregationResponse struct {
- Buckets []BucketBase `json:"buckets,omitempty"`
- Value interface{} `json:"value,omitempty"`
+ Buckets []BucketBase `json:"buckets,omitempty"`
+ Value interface{} `json:"value,omitempty"`
+ Interval string `json:"interval,omitempty"`
+}
+
+func (a *AggregationResponse) UnmarshalJSON(data []byte) error {
+ type alias struct {
+ Buckets json.RawMessage `json:"buckets,omitempty"`
+ Value interface{} `json:"value,omitempty"`
+ Interval string `json:"interval,omitempty"`
+ }
+
+ var aux alias
+ if err := json.Unmarshal(data, &aux); err != nil {
+ return err
+ }
+ a.Value = aux.Value
+ a.Interval = aux.Interval
+
+ buckets := bytes.TrimSpace(aux.Buckets)
+ if len(buckets) == 0 || bytes.Equal(buckets, []byte("null")) {
+ a.Buckets = nil
+ return nil
+ }
+
+ switch buckets[0] {
+ case '[':
+ return json.Unmarshal(buckets, &a.Buckets)
+ case '{':
+ keyedBuckets := map[string]BucketBase{}
+ if err := json.Unmarshal(buckets, &keyedBuckets); err != nil {
+ return err
+ }
+
+ keys := make([]string, 0, len(keyedBuckets))
+ for key := range keyedBuckets {
+ keys = append(keys, key)
+ }
+ sort.Slice(keys, func(i, j int) bool {
+ return compareBucketKeys(keys[i], keys[j])
+ })
+
+ a.Buckets = make([]BucketBase, 0, len(keys))
+ for _, key := range keys {
+ bucket := keyedBuckets[key]
+ if bucket == nil {
+ bucket = BucketBase{}
+ }
+ if _, ok := bucket["key"]; !ok {
+ bucket["key"] = key
+ }
+ a.Buckets = append(a.Buckets, bucket)
+ }
+ return nil
+ default:
+ return nil
+ }
+}
+
+func compareBucketKeys(left, right string) bool {
+ leftInt, leftErr := strconv.ParseInt(left, 10, 64)
+ rightInt, rightErr := strconv.ParseInt(right, 10, 64)
+ if leftErr == nil && rightErr == nil {
+ return leftInt < rightInt
+ }
+ return left < right
}
type ResponseBase struct {
@@ -235,6 +302,48 @@ type ErrorDetail struct {
Reason string `json:"reason,omitempty"`
}
+func (d *ErrorDetail) UnmarshalJSON(data []byte) error {
+ data = bytes.TrimSpace(data)
+ if len(data) == 0 || bytes.Equal(data, []byte("null")) {
+ return nil
+ }
+
+ if len(data) > 0 && data[0] == '"' {
+ return json.Unmarshal(data, &d.Reason)
+ }
+
+ type alias ErrorDetail
+ var aux alias
+ if err := json.Unmarshal(data, &aux); err != nil {
+ return err
+ }
+ *d = ErrorDetail(aux)
+ return nil
+}
+
+func (d *ErrorDetail) Message() string {
+ if d == nil {
+ return ""
+ }
+ if d.Reason != "" {
+ return d.Reason
+ }
+ if len(d.RootCause) > 0 {
+ var reasons []string
+ for _, cause := range d.RootCause {
+ if cause.Reason != "" {
+ reasons = append(reasons, cause.Reason)
+ } else if cause.Type != "" {
+ reasons = append(reasons, cause.Type)
+ }
+ }
+ if len(reasons) > 0 {
+ return strings.Join(reasons, "; ")
+ }
+ }
+ return d.Type
+}
+
type RootCause struct {
Type string `json:"type,omitempty"`
Reason string `json:"reason,omitempty"`
diff --git a/core/elastic/index_test.go b/core/elastic/index_test.go
index d36ae4d1c..399a9d6bc 100644
--- a/core/elastic/index_test.go
+++ b/core/elastic/index_test.go
@@ -25,8 +25,27 @@ package elastic
import (
"testing"
+
+ "github.com/segmentio/encoding/json"
)
+func TestAggregationResponseUnmarshalKeyedBuckets(t *testing.T) {
+ var agg AggregationResponse
+ err := json.Unmarshal([]byte(`{"buckets":{"0":{"doc_count":1740269},"1":{"doc_count":42}}}`), &agg)
+ if err != nil {
+ t.Fatalf("unexpected unmarshal error: %v", err)
+ }
+ if len(agg.Buckets) != 2 {
+ t.Fatalf("unexpected bucket count: %d", len(agg.Buckets))
+ }
+ if agg.Buckets[0]["key"] != "0" || agg.Buckets[0]["doc_count"] != float64(1740269) {
+ t.Fatalf("unexpected first bucket: %#v", agg.Buckets[0])
+ }
+ if agg.Buckets[1]["key"] != "1" || agg.Buckets[1]["doc_count"] != float64(42) {
+ t.Fatalf("unexpected second bucket: %#v", agg.Buckets[1])
+ }
+}
+
func TestIndexDocument_GetStringFieldFromSource(t *testing.T) {
tests := []struct {
name string
@@ -220,3 +239,31 @@ func TestIndexDocument_TryGetStringFieldFromSource(t *testing.T) {
})
}
}
+
+func TestErrorDetailUnmarshalJSONString(t *testing.T) {
+ var detail ErrorDetail
+ if err := json.Unmarshal([]byte(`"initializing"`), &detail); err != nil {
+ t.Fatalf("unexpected unmarshal error: %v", err)
+ }
+
+ if detail.Reason != "initializing" {
+ t.Fatalf("unexpected reason: %q", detail.Reason)
+ }
+ if detail.Message() != "initializing" {
+ t.Fatalf("unexpected message: %q", detail.Message())
+ }
+}
+
+func TestErrorDetailUnmarshalJSONObject(t *testing.T) {
+ var detail ErrorDetail
+ if err := json.Unmarshal([]byte(`{"type":"search_phase_execution_exception","reason":"all shards failed"}`), &detail); err != nil {
+ t.Fatalf("unexpected unmarshal error: %v", err)
+ }
+
+ if detail.Type != "search_phase_execution_exception" {
+ t.Fatalf("unexpected type: %q", detail.Type)
+ }
+ if detail.Message() != "all shards failed" {
+ t.Fatalf("unexpected message: %q", detail.Message())
+ }
+}
diff --git a/core/elastic/partition.go b/core/elastic/partition.go
index 63d474fb1..d18855642 100644
--- a/core/elastic/partition.go
+++ b/core/elastic/partition.go
@@ -32,6 +32,7 @@ import (
"fmt"
"math"
"net/http"
+ "sort"
"strconv"
"strings"
@@ -40,12 +41,14 @@ import (
)
type PartitionQuery struct {
- IndexName string `json:"index_name"`
- FieldType string `json:"field_type"`
- FieldName string `json:"field_name"`
- Step interface{} `json:"step"`
- Filter interface{} `json:"filter"`
- DocType string `json:"doc_type"`
+ IndexName string `json:"index_name"`
+ FieldType string `json:"field_type"`
+ FieldName string `json:"field_name"`
+ Strategy string `json:"strategy,omitempty"`
+ Step interface{} `json:"step,omitempty"`
+ PartitionCount int `json:"partition_count,omitempty"`
+ Filter interface{} `json:"filter"`
+ DocType string `json:"doc_type"`
}
type PartitionInfo struct {
@@ -54,6 +57,8 @@ type PartitionInfo struct {
End float64 `json:"end"`
Filter map[string]interface{} `json:"filter"`
Docs int64 `json:"docs"`
+ Label string `json:"label,omitempty"`
+ Values []string `json:"values,omitempty"`
Other bool
}
@@ -68,6 +73,11 @@ const (
PartitionByDate = "date"
PartitionByKeyword = "keyword"
PartitionByNumber = "number"
+
+ PartitionStrategyStep = "step"
+ PartitionStrategyQuantile = "quantile"
+ PartitionStrategyTerms = "terms"
+ PartitionStrategyHash = "hash"
)
func GetPartitions(q *PartitionQuery, client API) ([]PartitionInfo, error) {
@@ -100,32 +110,6 @@ func GetPartitions(q *PartitionQuery, client API) ([]PartitionInfo, error) {
switch q.FieldType {
case PartitionByDate, PartitionByNumber:
- var step float64
- if q.FieldType == PartitionByDate {
- if stepV, ok := q.Step.(string); !ok {
- return nil, fmt.Errorf("expect step value of string type since filedtype is %s", PartitionByDate)
- } else {
- du, err := util.ParseDuration(stepV)
- if err != nil {
- return nil, fmt.Errorf("parse step duration error: %w", err)
- }
- step = float64(du.Milliseconds())
- }
- } else {
- switch q.Step.(type) {
- case float64:
- step = q.Step.(float64)
- case string:
- v, err := strconv.Atoi(q.Step.(string))
- if err != nil {
- return nil, fmt.Errorf("convert step error: %w", err)
- }
- step = float64(v)
- default:
- return nil, fmt.Errorf("invalid parameter step: %v", q.Step)
- }
- }
-
result, err := getBoundValues(client, q.IndexName, q.FieldName, vFilter)
if err != nil {
return nil, err
@@ -138,23 +122,110 @@ func GetPartitions(q *PartitionQuery, client API) ([]PartitionInfo, error) {
var (
partitions []PartitionInfo
)
- partitions, err = getPartitionsByAgg(client, q.IndexName, q.FieldName, q.FieldType, step, vFilter)
- if err != nil {
- return nil, err
+
+ switch normalizePartitionStrategy(q.Strategy) {
+ case PartitionStrategyStep:
+ step, err := parsePartitionStep(q.FieldType, q.Step)
+ if err != nil {
+ return nil, err
+ }
+ partitions, err = getPartitionsByAgg(client, q.IndexName, q.FieldName, q.FieldType, step, vFilter)
+ if err != nil {
+ return nil, err
+ }
+ case PartitionStrategyQuantile:
+ partitions, err = getPartitionsByQuantile(client, q.IndexName, q.FieldName, q.FieldType, q.PartitionCount, result.Min, result.Max, vFilter)
+ if err != nil {
+ return nil, err
+ }
+ default:
+ return nil, fmt.Errorf("unsupported partition strategy: %s", q.Strategy)
}
+
if result.Null > 0 {
partitions = append(partitions, PartitionInfo{
Filter: result.NotExistsFilter,
Other: true,
+ Label: "Missing values",
Docs: result.Null,
})
}
return partitions, nil
+ case PartitionByKeyword:
+ var (
+ partitions []PartitionInfo
+ err error
+ )
+ switch normalizePartitionStrategy(q.Strategy) {
+ case PartitionStrategyTerms:
+ partitions, err = getPartitionsByTerms(client, q.IndexName, q.FieldName, q.PartitionCount, vFilter)
+ if err != nil {
+ return nil, err
+ }
+ case PartitionStrategyHash:
+ partitions, err = getPartitionsByHash(client, q.IndexName, q.FieldName, q.PartitionCount, vFilter)
+ if err != nil {
+ return nil, err
+ }
+ default:
+ return nil, fmt.Errorf("unsupported partition strategy: %s", q.Strategy)
+ }
+
+ missingPartition, err := getMissingPartition(client, q.IndexName, q.FieldName, vFilter)
+ if err != nil {
+ return nil, err
+ }
+ if missingPartition != nil {
+ partitions = append(partitions, *missingPartition)
+ }
+ return partitions, nil
default:
return nil, fmt.Errorf("unsupported field type: %s", q.FieldType)
}
}
+func normalizePartitionStrategy(strategy string) string {
+ switch strings.ToLower(strings.TrimSpace(strategy)) {
+ case "", PartitionStrategyStep:
+ return PartitionStrategyStep
+ case PartitionStrategyQuantile:
+ return PartitionStrategyQuantile
+ case PartitionStrategyTerms:
+ return PartitionStrategyTerms
+ case PartitionStrategyHash:
+ return PartitionStrategyHash
+ default:
+ return strings.ToLower(strings.TrimSpace(strategy))
+ }
+}
+
+func parsePartitionStep(fieldType string, stepValue interface{}) (float64, error) {
+ if fieldType == PartitionByDate {
+ stepV, ok := stepValue.(string)
+ if !ok {
+ return 0, fmt.Errorf("expect step value of string type since filedtype is %s", PartitionByDate)
+ }
+ du, err := util.ParseDuration(stepV)
+ if err != nil {
+ return 0, fmt.Errorf("parse step duration error: %w", err)
+ }
+ return float64(du.Milliseconds()), nil
+ }
+
+ switch stepValue.(type) {
+ case float64:
+ return stepValue.(float64), nil
+ case string:
+ v, err := strconv.Atoi(stepValue.(string))
+ if err != nil {
+ return 0, fmt.Errorf("convert step error: %w", err)
+ }
+ return float64(v), nil
+ default:
+ return 0, fmt.Errorf("invalid parameter step: %v", stepValue)
+ }
+}
+
func getPartitionsByAgg(client API, indexName string, fieldName, fieldType string, step float64, filter interface{}) ([]PartitionInfo, error) {
queryDsl := util.MapStr{
"size": 0,
@@ -182,7 +253,7 @@ func getPartitionsByAgg(client API, indexName string, fieldName, fieldType strin
if filter != nil {
queryDsl["query"] = filter
}
- res, err := client.SearchWithRawQueryDSL(indexName, util.MustToJSONBytes(queryDsl))
+ res, err := searchPartitionWithRawQueryDSL(client, indexName, queryDsl)
if err != nil {
return nil, err
}
@@ -217,13 +288,402 @@ func getPartitionsByAgg(client API, indexName string, fieldName, fieldType strin
Docs: int64(docCount),
Other: false,
}
- partition.Filter = buildPartitionFilter(min, max, fieldName, fieldType, filter)
+ partition.Filter = buildBoundedPartitionFilter(min, max, fieldName, fieldType, filter)
partitions = append(partitions, partition)
}
}
return partitions, nil
}
+func getPartitionsByQuantile(client API, indexName string, fieldName, fieldType string, partitionCount int, min, max float64, filter interface{}) ([]PartitionInfo, error) {
+ if partitionCount <= 0 {
+ return nil, fmt.Errorf("invalid parameter partition_count: %d", partitionCount)
+ }
+
+ boundaries, err := getQuantileBoundaries(client, indexName, fieldName, partitionCount, min, max, filter)
+ if err != nil {
+ return nil, err
+ }
+ partitions := buildQuantilePartitions(boundaries, fieldName, fieldType, filter)
+ if len(partitions) == 0 {
+ return nil, nil
+ }
+
+ counts, err := getPartitionDocCounts(client, indexName, partitions)
+ if err != nil {
+ return nil, err
+ }
+
+ filtered := make([]PartitionInfo, 0, len(partitions))
+ for i := range partitions {
+ partitions[i].Docs = counts[i]
+ if partitions[i].Docs <= 0 {
+ continue
+ }
+ filtered = append(filtered, partitions[i])
+ }
+ return filtered, nil
+}
+
+func getPartitionsByTerms(client API, indexName, fieldName string, partitionCount int, filter interface{}) ([]PartitionInfo, error) {
+ if partitionCount <= 0 {
+ return nil, fmt.Errorf("invalid parameter partition_count: %d", partitionCount)
+ }
+
+ queryDsl := util.MapStr{
+ "size": 0,
+ "aggs": util.MapStr{
+ "partitions": util.MapStr{
+ "terms": util.MapStr{
+ "field": fieldName,
+ "size": partitionCount,
+ },
+ },
+ },
+ }
+ if filter != nil {
+ queryDsl["query"] = filter
+ }
+
+ res, err := searchPartitionWithRawQueryDSL(client, indexName, queryDsl)
+ if err != nil {
+ return nil, err
+ }
+
+ var (
+ partitions []PartitionInfo
+ values []string
+ )
+ if partitionsAgg, ok := res.Aggregations["partitions"]; ok {
+ for idx, bucket := range partitionsAgg.Buckets {
+ value := fmt.Sprintf("%v", bucket["key"])
+ docCount := util.GetInt64Value(bucket["doc_count"])
+ if docCount <= 0 {
+ continue
+ }
+ values = append(values, value)
+ partitions = append(partitions, PartitionInfo{
+ Key: float64(idx),
+ Docs: docCount,
+ Label: value,
+ Values: []string{value},
+ Filter: buildExactTermPartitionFilter(value, fieldName, filter),
+ })
+ }
+ }
+
+ sumOtherDocCount, _ := jsonparser.GetInt(res.RawResult.Body, "aggregations", "partitions", "sum_other_doc_count")
+ if sumOtherDocCount > 0 {
+ partitions = append(partitions, PartitionInfo{
+ Key: float64(len(partitions)),
+ Docs: sumOtherDocCount,
+ Label: "Other terms",
+ Values: append([]string(nil), values...),
+ Filter: buildOtherTermsPartitionFilter(values, fieldName, filter),
+ Other: true,
+ })
+ }
+
+ return partitions, nil
+}
+
+func getPartitionsByHash(client API, indexName, fieldName string, partitionCount int, filter interface{}) ([]PartitionInfo, error) {
+ if partitionCount <= 0 {
+ return nil, fmt.Errorf("invalid parameter partition_count: %d", partitionCount)
+ }
+
+ partitions := make([]PartitionInfo, 0, partitionCount)
+ for idx := 0; idx < partitionCount; idx++ {
+ partitions = append(partitions, PartitionInfo{
+ Key: float64(idx),
+ Label: fmt.Sprintf("Hash %d/%d", idx+1, partitionCount),
+ Filter: buildHashPartitionFilter(idx, partitionCount, fieldName, filter),
+ })
+ }
+
+ counts, err := getHashPartitionDocCounts(client, indexName, fieldName, partitionCount, filter)
+ if err != nil {
+ return nil, err
+ }
+
+ filtered := make([]PartitionInfo, 0, len(partitions))
+ for idx := range partitions {
+ partitions[idx].Docs = counts[idx]
+ if partitions[idx].Docs <= 0 {
+ continue
+ }
+ filtered = append(filtered, partitions[idx])
+ }
+ return filtered, nil
+}
+
+func getHashPartitionDocCounts(client API, indexName, fieldName string, partitionCount int, filter interface{}) ([]int64, error) {
+ queryDsl := buildHashPartitionAggQuery(fieldName, partitionCount, filter)
+ res, err := searchPartitionWithRawQueryDSL(client, indexName, queryDsl)
+ if err != nil {
+ return nil, err
+ }
+ return extractHashPartitionDocCounts(res, partitionCount), nil
+}
+
+func buildHashPartitionAggQuery(fieldName string, partitionCount int, filter interface{}) util.MapStr {
+ fieldLiteral := buildPainlessStringLiteral(fieldName)
+ queryDsl := util.MapStr{
+ "size": 0,
+ "aggs": util.MapStr{
+ "partitions": util.MapStr{
+ "terms": util.MapStr{
+ "size": partitionCount,
+ "value_type": "long",
+ "script": util.MapStr{
+ "lang": "painless",
+ "source": fmt.Sprintf("if (doc[%s].size()==0 || doc[%s].value == '') return null; return (((doc[%s].value.hashCode() %% params.partition_count) + params.partition_count) %% params.partition_count);", fieldLiteral, fieldLiteral, fieldLiteral),
+ "params": util.MapStr{
+ "partition_count": partitionCount,
+ },
+ },
+ },
+ },
+ },
+ }
+ if filter != nil {
+ queryDsl["query"] = filter
+ }
+ return queryDsl
+}
+
+func extractHashPartitionDocCounts(res *SearchResponse, partitionCount int) []int64 {
+ counts := make([]int64, partitionCount)
+ if res == nil {
+ return counts
+ }
+ partitionsAgg, ok := res.Aggregations["partitions"]
+ if !ok {
+ return counts
+ }
+ for _, bucket := range partitionsAgg.Buckets {
+ bucketKey, ok := extractHashPartitionBucketKey(bucket["key"])
+ if !ok || bucketKey < 0 || bucketKey >= partitionCount {
+ continue
+ }
+ counts[bucketKey] = util.GetInt64Value(bucket["doc_count"])
+ }
+ return counts
+}
+
+func extractHashPartitionBucketKey(key interface{}) (int, bool) {
+ switch v := key.(type) {
+ case int:
+ return v, true
+ case int64:
+ return int(v), true
+ case int32:
+ return int(v), true
+ case uint:
+ return int(v), true
+ case uint64:
+ return int(v), true
+ case float64:
+ return int(v), true
+ case float32:
+ return int(v), true
+ case string:
+ parsed, err := strconv.Atoi(v)
+ if err != nil {
+ return 0, false
+ }
+ return parsed, true
+ default:
+ return 0, false
+ }
+}
+
+func getQuantileBoundaries(client API, indexName, fieldName string, partitionCount int, min, max float64, filter interface{}) ([]float64, error) {
+ percents := buildQuantilePercents(partitionCount)
+ if len(percents) == 0 {
+ return []float64{min, max}, nil
+ }
+
+ queryDsl := util.MapStr{
+ "size": 0,
+ "aggs": util.MapStr{
+ "partition_percentiles": util.MapStr{
+ "percentiles": util.MapStr{
+ "field": fieldName,
+ "percents": percents,
+ "keyed": false,
+ },
+ },
+ },
+ }
+ if filter != nil {
+ queryDsl["query"] = filter
+ }
+
+ res, err := searchPartitionWithRawQueryDSL(client, indexName, queryDsl)
+ if err != nil {
+ return nil, err
+ }
+
+ boundaries := make([]float64, 0, len(percents)+2)
+ boundaries = append(boundaries, min)
+ _, err = jsonparser.ArrayEach(res.RawResult.Body, func(value []byte, _ jsonparser.ValueType, _ int, err error) {
+ if err != nil {
+ return
+ }
+ boundary, parseErr := jsonparser.GetFloat(value, "value")
+ if parseErr != nil || math.IsNaN(boundary) || math.IsInf(boundary, 0) {
+ return
+ }
+ boundaries = append(boundaries, boundary)
+ }, "aggregations", "partition_percentiles", "values")
+ if err != nil {
+ return nil, err
+ }
+ boundaries = append(boundaries, max)
+ boundaries = dedupeSortedBoundaries(boundaries)
+ if len(boundaries) == 1 {
+ return []float64{boundaries[0], boundaries[0]}, nil
+ }
+ return boundaries, nil
+}
+
+func buildQuantilePercents(partitionCount int) []float64 {
+ if partitionCount <= 1 {
+ return nil
+ }
+ percents := make([]float64, 0, partitionCount-1)
+ for i := 1; i < partitionCount; i++ {
+ percents = append(percents, float64(i)*100/float64(partitionCount))
+ }
+ return percents
+}
+
+func dedupeSortedBoundaries(boundaries []float64) []float64 {
+ if len(boundaries) == 0 {
+ return nil
+ }
+ sort.Float64s(boundaries)
+ result := make([]float64, 0, len(boundaries))
+ for _, boundary := range boundaries {
+ if len(result) == 0 || !sameBoundary(result[len(result)-1], boundary) {
+ result = append(result, boundary)
+ }
+ }
+ return result
+}
+
+func sameBoundary(left, right float64) bool {
+ return math.Abs(left-right) <= 1e-9
+}
+
+func buildQuantilePartitions(boundaries []float64, fieldName, fieldType string, filter interface{}) []PartitionInfo {
+ if len(boundaries) < 2 {
+ return nil
+ }
+
+ partitions := make([]PartitionInfo, 0, len(boundaries)-1)
+ if len(boundaries) == 2 {
+ partitions = append(partitions, PartitionInfo{
+ Key: boundaries[1],
+ Start: boundaries[0],
+ End: boundaries[1],
+ Filter: buildOpenPartitionFilter(nil, nil, fieldName, fieldType, filter),
+ })
+ return partitions
+ }
+
+ for i := 1; i < len(boundaries); i++ {
+ lower, upper := boundaries[i-1], boundaries[i]
+ if sameBoundary(lower, upper) {
+ continue
+ }
+
+ var lowerRef, upperRef *float64
+ if i > 1 {
+ lowerRef = &lower
+ }
+ if i < len(boundaries)-1 {
+ upperRef = &upper
+ }
+
+ partitions = append(partitions, PartitionInfo{
+ Key: upper,
+ Start: lower,
+ End: upper,
+ Filter: buildOpenPartitionFilter(lowerRef, upperRef, fieldName, fieldType, filter),
+ })
+ }
+ return partitions
+}
+
+func getPartitionDocCounts(client API, indexName string, partitions []PartitionInfo) ([]int64, error) {
+ queryDsl := util.MapStr{
+ "size": 0,
+ "aggs": util.MapStr{
+ "partitions": util.MapStr{
+ "filters": util.MapStr{
+ "filters": buildPartitionFiltersMap(partitions),
+ },
+ },
+ },
+ }
+
+ res, err := searchPartitionWithRawQueryDSL(client, indexName, queryDsl)
+ if err != nil {
+ return nil, err
+ }
+
+ counts := make([]int64, 0, len(partitions))
+ for i := range partitions {
+ docCount, parseErr := jsonparser.GetInt(res.RawResult.Body, "aggregations", "partitions", "buckets", strconv.Itoa(i), "doc_count")
+ if parseErr != nil {
+ return nil, parseErr
+ }
+ counts = append(counts, docCount)
+ }
+ return counts, nil
+}
+
+func buildPartitionFiltersMap(partitions []PartitionInfo) util.MapStr {
+ filters := util.MapStr{}
+ for i, partition := range partitions {
+ filters[strconv.Itoa(i)] = partition.Filter
+ }
+ return filters
+}
+
+func getMissingPartition(client API, indexName, fieldName string, filter interface{}) (*PartitionInfo, error) {
+ queryDsl := util.MapStr{
+ "size": 0,
+ "aggs": util.MapStr{
+ "missing_field": util.MapStr{
+ "filter": buildMissingFieldCondition(fieldName),
+ },
+ },
+ }
+ if filter != nil {
+ queryDsl["query"] = filter
+ }
+
+ res, err := searchPartitionWithRawQueryDSL(client, indexName, queryDsl)
+ if err != nil {
+ return nil, err
+ }
+
+ docCount, err := jsonparser.GetInt(res.RawResult.Body, "aggregations", "missing_field", "doc_count")
+ if err != nil || docCount <= 0 {
+ return nil, err
+ }
+
+ return &PartitionInfo{
+ Docs: docCount,
+ Label: "Missing values",
+ Filter: buildMissingFieldFilter(fieldName, filter),
+ Other: true,
+ }, nil
+}
+
// NOTE: we assume GetPartitions returned sorted buckets from ES, if not, we need to manually sort source & target partitions by keys
// sourcePartitions & targetPartitions must've been generated with same bucket step & offset
func MergePartitions(sourcePartitions []PartitionInfo, targetPartitions []PartitionInfo, fieldName, fieldType string, filter interface{}) []PartitionInfo {
@@ -253,7 +713,7 @@ func MergePartitions(sourcePartitions []PartitionInfo, targetPartitions []Partit
Docs: util.MaxInt64(source.Docs, target.Docs),
Other: false,
}
- partition.Filter = buildPartitionFilter(partition.Start, partition.End, fieldName, fieldType, filter)
+ partition.Filter = buildBoundedPartitionFilter(partition.Start, partition.End, fieldName, fieldType, filter)
ret = append(ret, partition)
sourceIdx += 1
targetIdx += 1
@@ -267,12 +727,14 @@ func MergePartitions(sourcePartitions []PartitionInfo, targetPartitions []Partit
return ret
}
-func buildPartitionFilter(min, max float64, fieldName, fieldType string, filter interface{}) util.MapStr {
+func buildBoundedPartitionFilter(min, max float64, fieldName, fieldType string, filter interface{}) util.MapStr {
rv := util.MapStr{
"gte": min,
"lte": max,
}
if fieldType == PartitionByDate {
+ rv["gte"] = normalizeDateRangeBoundary(min, true, true)
+ rv["lte"] = normalizeDateRangeBoundary(max, false, true)
rv["format"] = "epoch_millis"
}
must := []interface{}{
@@ -290,7 +752,217 @@ func buildPartitionFilter(min, max float64, fieldName, fieldType string, filter
"must": must,
},
}
+}
+
+func buildOpenPartitionFilter(lower, upper *float64, fieldName, fieldType string, filter interface{}) util.MapStr {
+ rv := util.MapStr{}
+ if lower != nil {
+ rv["gt"] = *lower
+ }
+ if upper != nil {
+ rv["lte"] = *upper
+ }
+ if fieldType == PartitionByDate {
+ if lower != nil {
+ rv["gt"] = normalizeDateRangeBoundary(*lower, true, false)
+ }
+ if upper != nil {
+ rv["lte"] = normalizeDateRangeBoundary(*upper, false, true)
+ }
+ rv["format"] = "epoch_millis"
+ }
+ var condition interface{}
+ if len(rv) == 0 || (len(rv) == 1 && rv["format"] != nil) {
+ condition = util.MapStr{
+ "exists": util.MapStr{
+ "field": fieldName,
+ },
+ }
+ } else {
+ condition = util.MapStr{
+ "range": util.MapStr{
+ fieldName: rv,
+ },
+ }
+ }
+ must := []interface{}{condition}
+ if filter != nil {
+ must = append(must, filter)
+ }
+ return util.MapStr{
+ "bool": util.MapStr{
+ "must": must,
+ },
+ }
+
+}
+func normalizeDateRangeBoundary(value float64, lower, inclusive bool) int64 {
+ switch {
+ case lower && inclusive:
+ return int64(math.Ceil(value))
+ case lower && !inclusive:
+ return int64(math.Floor(value))
+ case !lower && inclusive:
+ return int64(math.Floor(value))
+ default:
+ return int64(math.Ceil(value))
+ }
+}
+
+func buildExactTermPartitionFilter(value, fieldName string, filter interface{}) util.MapStr {
+ return buildMustPartitionFilter([]interface{}{
+ util.MapStr{
+ "term": util.MapStr{
+ fieldName: util.MapStr{
+ "value": value,
+ },
+ },
+ },
+ }, filter)
+}
+
+func buildOtherTermsPartitionFilter(values []string, fieldName string, filter interface{}) util.MapStr {
+ boolFilter := util.MapStr{
+ "must": []interface{}{
+ util.MapStr{
+ "exists": util.MapStr{
+ "field": fieldName,
+ },
+ },
+ },
+ }
+ if filter != nil {
+ boolFilter["must"] = append(boolFilter["must"].([]interface{}), filter)
+ }
+ if len(values) > 0 {
+ boolFilter["must_not"] = []interface{}{
+ util.MapStr{
+ "terms": util.MapStr{
+ fieldName: values,
+ },
+ },
+ }
+ }
+ return util.MapStr{
+ "bool": boolFilter,
+ }
+}
+
+func buildHashPartitionFilter(partitionID, partitionCount int, fieldName string, filter interface{}) util.MapStr {
+ fieldLiteral := buildPainlessStringLiteral(fieldName)
+ return buildMustPartitionFilter([]interface{}{
+ util.MapStr{
+ "script": util.MapStr{
+ "script": util.MapStr{
+ "lang": "painless",
+ "source": fmt.Sprintf("doc[%s].size()!=0 && doc[%s].value != '' && (((doc[%s].value.hashCode() %% params.partition_count) + params.partition_count) %% params.partition_count) == params.partition_id", fieldLiteral, fieldLiteral, fieldLiteral),
+ "params": util.MapStr{
+ "partition_count": partitionCount,
+ "partition_id": partitionID,
+ },
+ },
+ },
+ },
+ }, filter)
+}
+
+func buildPainlessStringLiteral(value string) string {
+ replacer := strings.NewReplacer(`\`, `\\`, `'`, `\'`)
+ return "'" + replacer.Replace(value) + "'"
+}
+
+func searchPartitionWithRawQueryDSL(client API, indexName string, queryDsl util.MapStr) (*SearchResponse, error) {
+ res, err := client.SearchWithRawQueryDSL(indexName, util.MustToJSONBytes(queryDsl))
+ if err != nil {
+ return nil, err
+ }
+ if err := ensurePartitionSearchResponseOK(res); err != nil {
+ return nil, err
+ }
+ return res, nil
+}
+
+func ensurePartitionSearchResponseOK(res *SearchResponse) error {
+ if res == nil {
+ return errors.New("empty search response")
+ }
+ if res.StatusCode == 0 || res.StatusCode == http.StatusOK {
+ return nil
+ }
+ if res.RawResult != nil && len(res.RawResult.Body) > 0 {
+ for _, path := range [][]string{
+ {"error", "failed_shards", "[0]", "reason", "caused_by", "reason"},
+ {"error", "failed_shards", "[0]", "reason", "reason"},
+ {"error", "root_cause", "[0]", "reason"},
+ {"error", "reason"},
+ } {
+ if msg, ok := getJSONPathString(res.RawResult.Body, path...); ok && msg != "" {
+ return errors.New(msg)
+ }
+ }
+ }
+ if msg := res.Error.Message(); msg != "" {
+ return errors.New(msg)
+ }
+ if res.RawResult != nil && len(res.RawResult.Body) > 0 {
+ return errors.New(string(res.RawResult.Body))
+ }
+ return fmt.Errorf("unexpected search status: %d", res.StatusCode)
+}
+
+func getJSONPathString(data []byte, path ...string) (string, bool) {
+ v, err := jsonparser.GetString(data, path...)
+ if err != nil {
+ return "", false
+ }
+ return v, true
+}
+
+func buildMissingFieldCondition(fieldName string) util.MapStr {
+ return util.MapStr{
+ "bool": util.MapStr{
+ "should": []interface{}{
+ util.MapStr{
+ "bool": util.MapStr{
+ "must_not": []interface{}{
+ util.MapStr{
+ "exists": util.MapStr{
+ "field": fieldName,
+ },
+ },
+ },
+ },
+ },
+ util.MapStr{
+ "term": util.MapStr{
+ fieldName: util.MapStr{
+ "value": "",
+ },
+ },
+ },
+ },
+ "minimum_should_match": 1,
+ },
+ }
+}
+
+func buildMissingFieldFilter(fieldName string, filter interface{}) util.MapStr {
+ return buildMustPartitionFilter([]interface{}{
+ buildMissingFieldCondition(fieldName),
+ }, filter)
+}
+
+func buildMustPartitionFilter(mustClauses []interface{}, filter interface{}) util.MapStr {
+ must := append([]interface{}{}, mustClauses...)
+ if filter != nil {
+ must = append(must, filter)
+ }
+ return util.MapStr{
+ "bool": util.MapStr{
+ "must": must,
+ },
+ }
}
func getBoundValues(client API, indexName string, fieldName string, filter interface{}) (*BoundValuesResult, error) {
@@ -326,7 +998,7 @@ func getBoundValues(client API, indexName string, fieldName string, filter inter
if filter != nil {
queryDsl["query"] = filter
}
- res, err := client.SearchWithRawQueryDSL(indexName, util.MustToJSONBytes(queryDsl))
+ res, err := searchPartitionWithRawQueryDSL(client, indexName, queryDsl)
if err != nil {
return nil, err
}
diff --git a/core/elastic/partition_test.go b/core/elastic/partition_test.go
new file mode 100644
index 000000000..1a7795a6d
--- /dev/null
+++ b/core/elastic/partition_test.go
@@ -0,0 +1,298 @@
+package elastic
+
+import (
+ "net/http"
+ "reflect"
+ "strings"
+ "testing"
+
+ "infini.sh/framework/core/util"
+)
+
+func TestBuildQuantilePercents(t *testing.T) {
+ got := buildQuantilePercents(4)
+ want := []float64{25, 50, 75}
+ if !reflect.DeepEqual(got, want) {
+ t.Fatalf("unexpected percents: got %v want %v", got, want)
+ }
+}
+
+func TestBuildQuantilePartitionsCreatesOpenEdgeRanges(t *testing.T) {
+ partitions := buildQuantilePartitions([]float64{10, 20, 30}, "value", PartitionByNumber, nil)
+ if len(partitions) != 2 {
+ t.Fatalf("unexpected partition count: %d", len(partitions))
+ }
+
+ firstRange := getMustClause(t, partitions[0].Filter)["range"].(util.MapStr)["value"].(util.MapStr)
+ if _, ok := firstRange["gt"]; ok {
+ t.Fatalf("expected first partition to have no lower bound, got %v", firstRange)
+ }
+ if got := firstRange["lte"]; got != float64(20) {
+ t.Fatalf("unexpected first upper bound: %v", got)
+ }
+
+ secondRange := getMustClause(t, partitions[1].Filter)["range"].(util.MapStr)["value"].(util.MapStr)
+ if got := secondRange["gt"]; got != float64(20) {
+ t.Fatalf("unexpected second lower bound: %v", got)
+ }
+ if _, ok := secondRange["lte"]; ok {
+ t.Fatalf("expected last partition to have no upper bound, got %v", secondRange)
+ }
+}
+
+func TestBuildQuantilePartitionsSinglePartitionUsesExistsFilter(t *testing.T) {
+ partitions := buildQuantilePartitions([]float64{5, 5}, "value", PartitionByNumber, nil)
+ if len(partitions) != 1 {
+ t.Fatalf("unexpected partition count: %d", len(partitions))
+ }
+
+ clause := getMustClause(t, partitions[0].Filter)
+ exists, ok := clause["exists"].(util.MapStr)
+ if !ok {
+ t.Fatalf("expected exists clause, got %v", clause)
+ }
+ if exists["field"] != "value" {
+ t.Fatalf("unexpected exists field: %v", exists["field"])
+ }
+}
+
+func TestBuildOpenPartitionFilterPreservesDateFormat(t *testing.T) {
+ upper := 1000.0
+ filter := buildOpenPartitionFilter(nil, &upper, "ts", PartitionByDate, nil)
+ rangeFilter := getMustClause(t, filter)["range"].(util.MapStr)["ts"].(util.MapStr)
+ if got := rangeFilter["format"]; got != "epoch_millis" {
+ t.Fatalf("unexpected date format: %v", got)
+ }
+ if got := rangeFilter["lte"]; got != int64(1000) {
+ t.Fatalf("unexpected upper bound: %v", got)
+ }
+}
+
+func TestBuildOpenPartitionFilterRoundsDatePercentileBoundaries(t *testing.T) {
+ lower := 1779109187904.8455
+ upper := 1779109187999.999
+ filter := buildOpenPartitionFilter(&lower, &upper, "created_at", PartitionByDate, nil)
+ rangeFilter := getMustClause(t, filter)["range"].(util.MapStr)["created_at"].(util.MapStr)
+ if got := rangeFilter["gt"]; got != int64(1779109187904) {
+ t.Fatalf("unexpected lower bound: %v", got)
+ }
+ if got := rangeFilter["lte"]; got != int64(1779109187999) {
+ t.Fatalf("unexpected upper bound: %v", got)
+ }
+}
+
+func TestBuildBoundedPartitionFilterRoundsDateBoundaries(t *testing.T) {
+ filter := buildBoundedPartitionFilter(1779109187904.1, 1779109187999.9, "created_at", PartitionByDate, nil)
+ rangeFilter := getMustClause(t, filter)["range"].(util.MapStr)["created_at"].(util.MapStr)
+ if got := rangeFilter["gte"]; got != int64(1779109187905) {
+ t.Fatalf("unexpected lower bound: %v", got)
+ }
+ if got := rangeFilter["lte"]; got != int64(1779109187999) {
+ t.Fatalf("unexpected upper bound: %v", got)
+ }
+}
+
+func TestBuildExactTermPartitionFilter(t *testing.T) {
+ filter := buildExactTermPartitionFilter("pmid-1", "pmid.keyword", nil)
+ termFilter := getMustClause(t, filter)["term"].(util.MapStr)["pmid.keyword"].(util.MapStr)
+ if got := termFilter["value"]; got != "pmid-1" {
+ t.Fatalf("unexpected term value: %v", got)
+ }
+}
+
+func TestBuildOtherTermsPartitionFilter(t *testing.T) {
+ filter := buildOtherTermsPartitionFilter([]string{"a", "b"}, "pmid.keyword", nil)
+ boolFilter := filter["bool"].(util.MapStr)
+ mustNot := boolFilter["must_not"].([]interface{})
+ termsFilter := mustNot[0].(util.MapStr)["terms"].(util.MapStr)
+ values := termsFilter["pmid.keyword"].([]string)
+ if !reflect.DeepEqual(values, []string{"a", "b"}) {
+ t.Fatalf("unexpected excluded values: %v", values)
+ }
+}
+
+func TestBuildHashPartitionFilter(t *testing.T) {
+ filter := buildHashPartitionFilter(1, 8, "pmid.keyword", nil)
+ scriptFilter := getMustClause(t, filter)["script"].(util.MapStr)["script"].(util.MapStr)
+ if scriptFilter["lang"] != "painless" {
+ t.Fatalf("unexpected script language: %v", scriptFilter["lang"])
+ }
+ source, ok := scriptFilter["source"].(string)
+ if !ok {
+ t.Fatalf("unexpected script source: %T", scriptFilter["source"])
+ }
+ if !strings.Contains(source, "doc['pmid.keyword']") {
+ t.Fatalf("unexpected script source: %s", source)
+ }
+ if !strings.Contains(source, "value != ''") {
+ t.Fatalf("expected empty strings to be excluded from hash partition, got %s", source)
+ }
+ if strings.Contains(source, "Math.floorMod") {
+ t.Fatalf("unexpected script source: %s", source)
+ }
+ params := scriptFilter["params"].(util.MapStr)
+ if params["partition_count"] != 8 || params["partition_id"] != 1 {
+ t.Fatalf("unexpected script params: %v", params)
+ }
+ if _, ok := params["field"]; ok {
+ t.Fatalf("field should not be passed as a script param: %v", params)
+ }
+}
+
+func TestBuildHashPartitionAggQueryAppliesOuterFilter(t *testing.T) {
+ query := buildHashPartitionAggQuery("pmid.keyword", 8, util.MapStr{
+ "term": util.MapStr{
+ "env": util.MapStr{"value": "prod"},
+ },
+ })
+
+ if !reflect.DeepEqual(query["query"], util.MapStr{
+ "term": util.MapStr{
+ "env": util.MapStr{"value": "prod"},
+ },
+ }) {
+ t.Fatalf("expected outer filter to be applied at top-level query, got %v", query["query"])
+ }
+
+ termsAgg := query["aggs"].(util.MapStr)["partitions"].(util.MapStr)["terms"].(util.MapStr)
+ if got := termsAgg["size"]; got != 8 {
+ t.Fatalf("unexpected partition size: %v", got)
+ }
+ if got := termsAgg["value_type"]; got != "long" {
+ t.Fatalf("unexpected value_type: %v", got)
+ }
+ script := termsAgg["script"].(util.MapStr)
+ source, ok := script["source"].(string)
+ if !ok {
+ t.Fatalf("unexpected script source type: %T", script["source"])
+ }
+ if !strings.Contains(source, "return null") {
+ t.Fatalf("expected missing values to be skipped in hash aggregation, got %s", source)
+ }
+ if !strings.Contains(source, "value == ''") {
+ t.Fatalf("expected empty strings to be excluded in hash aggregation, got %s", source)
+ }
+ params := script["params"].(util.MapStr)
+ if got := params["partition_count"]; got != 8 {
+ t.Fatalf("unexpected partition_count: %v", got)
+ }
+}
+
+func TestExtractHashPartitionDocCountsMapsByBucketKey(t *testing.T) {
+ counts := extractHashPartitionDocCounts(&SearchResponse{
+ Aggregations: map[string]AggregationResponse{
+ "partitions": {
+ Buckets: []BucketBase{
+ {"key": float64(5), "doc_count": float64(12)},
+ {"key": "1", "doc_count": float64(7)},
+ {"key": float64(99), "doc_count": float64(3)},
+ },
+ },
+ },
+ }, 8)
+
+ expected := []int64{0, 7, 0, 0, 0, 12, 0, 0}
+ if !reflect.DeepEqual(counts, expected) {
+ t.Fatalf("unexpected hash counts: got %v want %v", counts, expected)
+ }
+}
+
+func TestBuildMissingFieldConditionIncludesEmptyString(t *testing.T) {
+ filter := buildMissingFieldCondition("pmid.keyword")
+ boolFilter, ok := filter["bool"].(util.MapStr)
+ if !ok {
+ t.Fatalf("expected bool filter, got %v", filter)
+ }
+ if got := boolFilter["minimum_should_match"]; got != 1 {
+ t.Fatalf("unexpected minimum_should_match: %v", got)
+ }
+ should, ok := boolFilter["should"].([]interface{})
+ if !ok || len(should) != 2 {
+ t.Fatalf("expected two should clauses, got %v", boolFilter["should"])
+ }
+ termFilter := should[1].(util.MapStr)["term"].(util.MapStr)["pmid.keyword"].(util.MapStr)
+ if got := termFilter["value"]; got != "" {
+ t.Fatalf("unexpected empty-string term filter: %v", termFilter)
+ }
+}
+
+func TestBuildMissingFieldFilterPreservesOuterFilter(t *testing.T) {
+ filter := buildMissingFieldFilter("pmid.keyword", util.MapStr{
+ "term": util.MapStr{
+ "env": util.MapStr{"value": "prod"},
+ },
+ })
+ boolFilter, ok := filter["bool"].(util.MapStr)
+ if !ok {
+ t.Fatalf("expected bool filter, got %v", filter)
+ }
+ must, ok := boolFilter["must"].([]interface{})
+ if !ok || len(must) != 2 {
+ t.Fatalf("expected two must clauses, got %v", boolFilter["must"])
+ }
+ innerBool, ok := must[0].(util.MapStr)["bool"].(util.MapStr)
+ if !ok {
+ t.Fatalf("expected wrapped missing bool filter, got %v", must[0])
+ }
+ if got := innerBool["minimum_should_match"]; got != 1 {
+ t.Fatalf("unexpected minimum_should_match: %v", got)
+ }
+}
+
+func TestBuildPainlessStringLiteralEscapesSingleQuote(t *testing.T) {
+ got := buildPainlessStringLiteral("foo'bar")
+ if got != `'foo\'bar'` {
+ t.Fatalf("unexpected painless string literal: %s", got)
+ }
+}
+
+func TestEnsurePartitionSearchResponseOKReturnsBackendReason(t *testing.T) {
+ err := ensurePartitionSearchResponseOK(&SearchResponse{
+ ResponseBase: ResponseBase{
+ StatusCode: http.StatusInternalServerError,
+ RawResult: &util.Result{
+ Body: []byte(`{"error":{"reason":"runtime script failure"},"status":500}`),
+ },
+ InternalError: InternalError{
+ Error: &ErrorDetail{
+ Reason: "runtime script failure",
+ },
+ Status: http.StatusInternalServerError,
+ },
+ },
+ })
+ if err == nil || err.Error() != "runtime script failure" {
+ t.Fatalf("unexpected error: %v", err)
+ }
+}
+
+func TestEnsurePartitionSearchResponseOKReturnsCausedByReason(t *testing.T) {
+ err := ensurePartitionSearchResponseOK(&SearchResponse{
+ ResponseBase: ResponseBase{
+ StatusCode: http.StatusBadRequest,
+ RawResult: &util.Result{
+ Body: []byte(`{"error":{"root_cause":[{"reason":"compile error"}],"failed_shards":[{"reason":{"reason":"compile error","caused_by":{"reason":"static method [java.lang.Math, floorMod/2] not found"}}}],"reason":"all shards failed"},"status":400}`),
+ },
+ },
+ })
+ if err == nil || err.Error() != "static method [java.lang.Math, floorMod/2] not found" {
+ t.Fatalf("unexpected error: %v", err)
+ }
+}
+
+func getMustClause(t *testing.T, filter util.MapStr) util.MapStr {
+ t.Helper()
+ boolFilter, ok := filter["bool"].(util.MapStr)
+ if !ok {
+ t.Fatalf("expected bool filter, got %v", filter)
+ }
+ must, ok := boolFilter["must"].([]interface{})
+ if !ok || len(must) == 0 {
+ t.Fatalf("expected must clauses, got %v", boolFilter["must"])
+ }
+ clause, ok := must[0].(util.MapStr)
+ if !ok {
+ t.Fatalf("expected util.MapStr clause, got %T", must[0])
+ }
+ return clause
+}
diff --git a/core/env/env.go b/core/env/env.go
index 45b61ff8e..68b5df3f0 100755
--- a/core/env/env.go
+++ b/core/env/env.go
@@ -294,9 +294,13 @@ func GetDefaultSystemConfig() config.SystemConfig {
},
Security: config.WebSecurityConfig{
Enabled: true,
- Authentication: config.AuthenticationConfig{Native: config.RealmConfig{
- Enabled: false,
- },
+ Authentication: config.AuthenticationConfig{
+ Native: config.RealmConfig{
+ Enabled: false,
+ },
+ AccessToken: config.AccessTokenConfig{
+ Enabled: true,
+ },
},
},
WebsocketConfig: config.WebsocketConfig{
diff --git a/core/env/env_test.go b/core/env/env_test.go
index 23e00d752..74afd9f2e 100644
--- a/core/env/env_test.go
+++ b/core/env/env_test.go
@@ -31,6 +31,14 @@ import (
"infini.sh/framework/core/config"
)
+func TestGetDefaultSystemConfigEnablesAccessTokenAPI(t *testing.T) {
+ cfg := GetDefaultSystemConfig()
+
+ if !cfg.WebAppConfig.Security.Authentication.AccessToken.Enabled {
+ t.Fatal("expected access token api to be enabled by default")
+ }
+}
+
func TestParseConfigSection_NilConfig(t *testing.T) {
var out struct{ Foo string }
exist, err := ParseConfigSection(nil, "anykey", &out)
diff --git a/core/env/http_client.go b/core/env/http_client.go
index f5ba79253..e154e2ac5 100644
--- a/core/env/http_client.go
+++ b/core/env/http_client.go
@@ -41,6 +41,30 @@ func (env *Env) GetHTTPClientConfig(name, endpoint string) *config.HTTPClientCon
TLSConfig: config.TLSConfig{SkipDomainVerify: true, TLSInsecureSkipVerify: true},
}
}
+ if name == "configs" && ((!ok && !isZeroTLSConfig(env.SystemConfig.Configs.TLSConfig)) || isZeroTLSConfig(clientCfg.TLSConfig)) {
+ clientCfg.TLSConfig = env.SystemConfig.Configs.TLSConfig
+ }
//TODO support client config per endpoint
return &clientCfg
}
+
+func isZeroTLSConfig(cfg config.TLSConfig) bool {
+ return !cfg.TLSEnabled &&
+ cfg.TLSCertFile == "" &&
+ cfg.TLSCertPassword == "" &&
+ cfg.TLSKeyFile == "" &&
+ cfg.TLSCACertFile == "" &&
+ !cfg.TLSInsecureSkipVerify &&
+ cfg.DefaultDomain == "" &&
+ !cfg.SkipDomainVerify &&
+ cfg.ClientSessionCacheSize == 0 &&
+ !cfg.TLSBypassMalformedCert &&
+ !cfg.AutoIssue.Enabled &&
+ cfg.AutoIssue.Email == "" &&
+ cfg.AutoIssue.Path == "" &&
+ !cfg.AutoIssue.IncludeDefaultDomain &&
+ !cfg.AutoIssue.SkipInvalidDomain &&
+ len(cfg.AutoIssue.Domains) == 0 &&
+ cfg.AutoIssue.Provider.TencentDNS.SecretID == "" &&
+ cfg.AutoIssue.Provider.TencentDNS.SecretKey == ""
+}
diff --git a/core/env/http_client_test.go b/core/env/http_client_test.go
new file mode 100644
index 000000000..1fb96ec11
--- /dev/null
+++ b/core/env/http_client_test.go
@@ -0,0 +1,63 @@
+package env
+
+import (
+ "testing"
+
+ "infini.sh/framework/core/config"
+)
+
+func TestGetHTTPClientConfigFallsBackToConfigsTLS(t *testing.T) {
+ env := &Env{
+ SystemConfig: &config.SystemConfig{
+ Configs: config.ConfigsConfig{
+ TLSConfig: config.TLSConfig{
+ TLSEnabled: true,
+ TLSCertFile: "config/client.crt",
+ TLSKeyFile: "config/client.key",
+ TLSCACertFile: "config/ca.crt",
+ TLSInsecureSkipVerify: false,
+ SkipDomainVerify: true,
+ ClientSessionCacheSize: 64,
+ },
+ },
+ },
+ }
+
+ cfg := env.GetHTTPClientConfig("configs", "")
+ if cfg.TLSConfig.TLSCertFile != "config/client.crt" {
+ t.Fatalf("expected configs tls cert_file fallback, got %q", cfg.TLSConfig.TLSCertFile)
+ }
+ if cfg.TLSConfig.TLSKeyFile != "config/client.key" {
+ t.Fatalf("expected configs tls key_file fallback, got %q", cfg.TLSConfig.TLSKeyFile)
+ }
+ if cfg.TLSConfig.TLSCACertFile != "config/ca.crt" {
+ t.Fatalf("expected configs tls ca_file fallback, got %q", cfg.TLSConfig.TLSCACertFile)
+ }
+ if !cfg.TLSConfig.SkipDomainVerify {
+ t.Fatal("expected configs tls skip_domain_verify fallback")
+ }
+}
+
+func TestGetHTTPClientConfigKeepsExplicitConfigsClientTLS(t *testing.T) {
+ env := &Env{
+ SystemConfig: &config.SystemConfig{
+ Configs: config.ConfigsConfig{
+ TLSConfig: config.TLSConfig{
+ TLSCertFile: "config/client.crt",
+ },
+ },
+ HTTPClientConfig: map[string]config.HTTPClientConfig{
+ "configs": {
+ TLSConfig: config.TLSConfig{
+ TLSCertFile: "override/client.crt",
+ },
+ },
+ },
+ },
+ }
+
+ cfg := env.GetHTTPClientConfig("configs", "")
+ if cfg.TLSConfig.TLSCertFile != "override/client.crt" {
+ t.Fatalf("expected explicit configs client tls to win, got %q", cfg.TLSConfig.TLSCertFile)
+ }
+}
diff --git a/core/kv/kv.go b/core/kv/kv.go
index 374e1893d..fc0695884 100755
--- a/core/kv/kv.go
+++ b/core/kv/kv.go
@@ -30,6 +30,7 @@ package kv
import (
log "github.com/cihub/seelog"
"infini.sh/framework/core/errors"
+ "time"
)
type KVStore interface {
@@ -42,8 +43,10 @@ type KVStore interface {
GetCompressedValue(bucket string, key []byte) ([]byte, error)
AddValueCompress(bucket string, key []byte, value []byte) error
+ AddValueCompressWithTTL(bucket string, key []byte, value []byte, ttl time.Duration) error
AddValue(bucket string, key []byte, value []byte) error
+ AddValueWithTTL(bucket string, key []byte, value []byte, ttl time.Duration) error
ExistsKey(bucket string, key []byte) (bool, error)
@@ -54,6 +57,14 @@ type KVStore interface {
var handler KVStore
+func HasStore(name string) bool {
+ if stores == nil {
+ return false
+ }
+ _, ok := stores[name]
+ return ok
+}
+
func getKVHandler() KVStore {
if handler == nil {
@@ -74,10 +85,18 @@ func AddValueCompress(bucket string, key []byte, value []byte) error {
return getKVHandler().AddValueCompress(bucket, key, value)
}
+func AddValueCompressWithTTL(bucket string, key []byte, value []byte, ttl time.Duration) error {
+ return getKVHandler().AddValueCompressWithTTL(bucket, key, value, ttl)
+}
+
func AddValue(bucket string, key []byte, value []byte) error {
return getKVHandler().AddValue(bucket, key, value)
}
+func AddValueWithTTL(bucket string, key []byte, value []byte, ttl time.Duration) error {
+ return getKVHandler().AddValueWithTTL(bucket, key, value, ttl)
+}
+
func ExistsKey(bucket string, key []byte) (bool, error) {
return getKVHandler().ExistsKey(bucket, key)
}
diff --git a/core/kv/kv_test.go b/core/kv/kv_test.go
new file mode 100644
index 000000000..0102e97c1
--- /dev/null
+++ b/core/kv/kv_test.go
@@ -0,0 +1,25 @@
+package kv
+
+import "testing"
+
+func TestHasStore(t *testing.T) {
+ previousHandler := handler
+ previousStores := stores
+ defer func() {
+ handler = previousHandler
+ stores = previousStores
+ }()
+
+ handler = nil
+ stores = nil
+
+ if HasStore("elastic") {
+ t.Fatal("expected store lookup to be false before registration")
+ }
+
+ Register("elastic", nil)
+
+ if !HasStore("elastic") {
+ t.Fatal("expected store lookup to be true after registration")
+ }
+}
diff --git a/core/model/const.go b/core/model/const.go
new file mode 100644
index 000000000..7a9cdca63
--- /dev/null
+++ b/core/model/const.go
@@ -0,0 +1,29 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package model
+
+const (
+ CredentialIDSystemKey = "credential_id"
+ API_TOKEN = "X-API-TOKEN"
+)
diff --git a/core/model/instance.go b/core/model/instance.go
index f9c60b44b..e5064ed7a 100644
--- a/core/model/instance.go
+++ b/core/model/instance.go
@@ -34,6 +34,7 @@ import (
"time"
log "github.com/cihub/seelog"
+ "infini.sh/framework/core/config"
"infini.sh/framework/core/env"
"infini.sh/framework/core/global"
"infini.sh/framework/core/host"
@@ -54,7 +55,11 @@ type Instance struct {
//application information
Application env.Application `json:"application,omitempty" elastic_mapping:"application: { type: object }"`
- BasicAuth *BasicAuth `config:"basic_auth" json:"basic_auth,omitempty" elastic_mapping:"basic_auth:{type:object}"`
+ BasicAuth *BasicAuth `config:"basic_auth" json:"basic_auth,omitempty" elastic_mapping:"basic_auth:{type:object}"`
+ AccessToken *Token `config:"access_token" json:"access_token,omitempty" elastic_mapping:"access_token:{type:object}"`
+
+ ManagerCredentialID string `json:"manager_credential_id,omitempty" elastic_mapping:"manager_credential_id:{type:keyword}"`
+ AccessCredentialID string `json:"access_credential_id,omitempty" elastic_mapping:"access_credential_id:{type:keyword}"`
Labels map[string]string `json:"labels,omitempty" elastic_mapping:"labels:{type:object}"`
Tags []string `json:"tags,omitempty"`
@@ -126,6 +131,33 @@ func (inst *Instance) GetVersion() (map[string]interface{}, error) {
return nil, fmt.Errorf("unknow agent version")
}
+func resolveManagedInstanceEndpoint(apiConfig config.APIConfig, webConfig config.WebAppConfig) string {
+ if apiConfig.Enabled {
+ return apiConfig.GetEndpoint()
+ }
+ if webConfig.Enabled {
+ return webConfig.GetEndpoint()
+ }
+ return apiConfig.GetEndpoint()
+}
+
+func buildManagedInstanceServices(apiConfig config.APIConfig, webConfig config.WebAppConfig) []ServiceInfo {
+ services := []ServiceInfo{}
+ if apiConfig.Enabled {
+ services = append(services, ServiceInfo{
+ Name: "api",
+ Endpoint: apiConfig.GetEndpoint(),
+ })
+ }
+ if webConfig.Enabled {
+ services = append(services, ServiceInfo{
+ Name: "web",
+ Endpoint: webConfig.GetEndpoint(),
+ })
+ }
+ return services
+}
+
func GetInstanceInfo() Instance {
instance := Instance{}
instance.ID = global.Env().SystemConfig.NodeConfig.ID
@@ -137,7 +169,8 @@ func GetInstanceInfo() Instance {
_, publicIP, _, _ := util.GetPublishNetworkDeviceInfo(global.Env().SystemConfig.NodeConfig.MajorIpPattern)
- instance.Endpoint = global.Env().SystemConfig.APIConfig.GetEndpoint()
+ instance.Endpoint = resolveManagedInstanceEndpoint(global.Env().SystemConfig.APIConfig, global.Env().SystemConfig.WebAppConfig)
+ instance.Services = buildManagedInstanceServices(global.Env().SystemConfig.APIConfig, global.Env().SystemConfig.WebAppConfig)
ips := util.GetLocalIPs()
if len(ips) > 0 {
diff --git a/core/model/instance_test.go b/core/model/instance_test.go
new file mode 100644
index 000000000..9713684de
--- /dev/null
+++ b/core/model/instance_test.go
@@ -0,0 +1,51 @@
+package model
+
+import (
+ "testing"
+
+ "infini.sh/framework/core/config"
+)
+
+func TestResolveManagedInstanceEndpoint(t *testing.T) {
+ t.Run("prefer api endpoint when api is enabled", func(t *testing.T) {
+ apiConfig := config.APIConfig{Enabled: true}
+ apiConfig.NetworkConfig.Publish = "127.0.0.1:2900"
+ webConfig := config.WebAppConfig{Enabled: true}
+ webConfig.NetworkConfig.Publish = "127.0.0.1:8080"
+
+ endpoint := resolveManagedInstanceEndpoint(apiConfig, webConfig)
+ if endpoint != "http://127.0.0.1:2900" {
+ t.Fatalf("unexpected endpoint: %s", endpoint)
+ }
+ })
+
+ t.Run("fallback to web endpoint when api is disabled", func(t *testing.T) {
+ apiConfig := config.APIConfig{Enabled: false}
+ apiConfig.NetworkConfig.Publish = "127.0.0.1:2900"
+ webConfig := config.WebAppConfig{Enabled: true}
+ webConfig.NetworkConfig.Publish = "127.0.0.1:8080"
+
+ endpoint := resolveManagedInstanceEndpoint(apiConfig, webConfig)
+ if endpoint != "http://127.0.0.1:8080" {
+ t.Fatalf("unexpected endpoint: %s", endpoint)
+ }
+ })
+}
+
+func TestBuildManagedInstanceServices(t *testing.T) {
+ apiConfig := config.APIConfig{Enabled: true}
+ apiConfig.NetworkConfig.Publish = "127.0.0.1:2900"
+ webConfig := config.WebAppConfig{Enabled: true}
+ webConfig.NetworkConfig.Publish = "127.0.0.1:8080"
+
+ services := buildManagedInstanceServices(apiConfig, webConfig)
+ if len(services) != 2 {
+ t.Fatalf("unexpected service count: %#v", services)
+ }
+ if services[0].Name != "api" || services[0].Endpoint != "http://127.0.0.1:2900" {
+ t.Fatalf("unexpected api service: %#v", services[0])
+ }
+ if services[1].Name != "web" || services[1].Endpoint != "http://127.0.0.1:8080" {
+ t.Fatalf("unexpected web service: %#v", services[1])
+ }
+}
diff --git a/core/model/token.go b/core/model/token.go
new file mode 100644
index 000000000..25fc4c74d
--- /dev/null
+++ b/core/model/token.go
@@ -0,0 +1,32 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+/* Copyright © INFINI LTD. All rights reserved.
+ * Web: https://infinilabs.com
+ * Email: hello#infini.ltd */
+
+package model
+
+type Token struct {
+ Value string `json:"value,omitempty" config:"value"`
+}
diff --git a/core/orm/orm_test.go b/core/orm/orm_test.go
index cfba7c8ca..082e2ef4d 100644
--- a/core/orm/orm_test.go
+++ b/core/orm/orm_test.go
@@ -158,6 +158,28 @@ func TestSetFieldTimeValue(t *testing.T) {
}
+func TestHasAdapter(t *testing.T) {
+ previousHandler := handler
+ previousAdapters := adapters
+ defer func() {
+ handler = previousHandler
+ adapters = previousAdapters
+ }()
+
+ handler = nil
+ adapters = nil
+
+ if HasAdapter("elastic") {
+ t.Fatal("expected adapter lookup to be false before registration")
+ }
+
+ Register("elastic", nil)
+
+ if !HasAdapter("elastic") {
+ t.Fatal("expected adapter lookup to be true after registration")
+ }
+}
+
//func TestSetFieldTimeValue1(t *testing.T) {
// t1:=time.Now()
// a:=struct {
diff --git a/core/orm/registry.go b/core/orm/registry.go
index 292b42e36..bf98277d2 100644
--- a/core/orm/registry.go
+++ b/core/orm/registry.go
@@ -10,6 +10,11 @@ import (
var registeredSchemas = []util.KeyValue{}
+func schemaRegistrationKey(t interface{}) string {
+ pkg, typeName := util.GetTypeAndPackageName(t, true)
+ return pkg + "-" + typeName
+}
+
func MustRegisterSchemaWithIndexName(t interface{}, index string) {
err := RegisterSchemaWithIndexName(t, index)
if err != nil {
@@ -18,6 +23,17 @@ func MustRegisterSchemaWithIndexName(t interface{}, index string) {
}
func RegisterSchemaWithIndexName(t interface{}, index string) error {
+ newKey := schemaRegistrationKey(t)
+ for _, registered := range registeredSchemas {
+ if registered.Key != index {
+ continue
+ }
+ existingKey := schemaRegistrationKey(registered.Payload)
+ if existingKey == newKey {
+ return nil
+ }
+ return errors.Errorf("schema index [%s] already registered by [%s]", index, existingKey)
+ }
registeredSchemas = append(registeredSchemas, util.KeyValue{Key: index, Payload: t})
return nil
}
@@ -35,6 +51,18 @@ func InitSchema() error {
var handler ORM
+func HasHandler() bool {
+ return handler != nil
+}
+
+func HasAdapter(name string) bool {
+ if adapters == nil {
+ return false
+ }
+ _, ok := adapters[name]
+ return ok
+}
+
func getHandler() ORM {
if handler == nil {
panic(errors.New("ORM handler is not registered"))
diff --git a/core/orm/registry_test.go b/core/orm/registry_test.go
new file mode 100644
index 000000000..b74dc34dd
--- /dev/null
+++ b/core/orm/registry_test.go
@@ -0,0 +1,39 @@
+package orm
+
+import "testing"
+
+type testSchemaAlpha struct{}
+type testSchemaBeta struct{}
+
+func TestRegisterSchemaWithIndexNameDeduplicatesSameSchema(t *testing.T) {
+ original := registeredSchemas
+ registeredSchemas = nil
+ t.Cleanup(func() {
+ registeredSchemas = original
+ })
+
+ if err := RegisterSchemaWithIndexName(testSchemaAlpha{}, "test-index"); err != nil {
+ t.Fatalf("expected first registration to succeed, got %v", err)
+ }
+ if err := RegisterSchemaWithIndexName(&testSchemaAlpha{}, "test-index"); err != nil {
+ t.Fatalf("expected duplicate registration to be ignored, got %v", err)
+ }
+ if len(registeredSchemas) != 1 {
+ t.Fatalf("expected exactly one registered schema, got %d", len(registeredSchemas))
+ }
+}
+
+func TestRegisterSchemaWithIndexNameRejectsDifferentSchemaForSameIndex(t *testing.T) {
+ original := registeredSchemas
+ registeredSchemas = nil
+ t.Cleanup(func() {
+ registeredSchemas = original
+ })
+
+ if err := RegisterSchemaWithIndexName(testSchemaAlpha{}, "test-index"); err != nil {
+ t.Fatalf("expected first registration to succeed, got %v", err)
+ }
+ if err := RegisterSchemaWithIndexName(testSchemaBeta{}, "test-index"); err == nil {
+ t.Fatal("expected conflicting registration to fail")
+ }
+}
diff --git a/core/pipeline/context.go b/core/pipeline/context.go
index 97fe18af8..4c8c6e88d 100755
--- a/core/pipeline/context.go
+++ b/core/pipeline/context.go
@@ -299,6 +299,20 @@ func (ctx *Context) Errors() []error {
return ctx.processErrs
}
+func (ctx *Context) GetResultState() RunningState {
+ ctx.stateLock.Lock()
+ defer ctx.stateLock.Unlock()
+
+ return ctx.getResultStateLocked()
+}
+
+func (ctx *Context) GetResultError() string {
+ ctx.stateLock.Lock()
+ defer ctx.stateLock.Unlock()
+
+ return formatPipelineResultError(ctx.exitErr, ctx.processErrs)
+}
+
// Pause suspends the goroutine that is running this pipeline.
func (ctx *Context) Pause() {
ctx.stateLock.Lock()
@@ -378,6 +392,30 @@ func (ctx *Context) setRunningState(newState RunningState) {
}
}
+func (ctx *Context) getResultStateLocked() RunningState {
+ switch ctx.runningState {
+ case FINISHED, FAILED:
+ return ctx.runningState
+ case STOPPED:
+ if ctx.endTime == nil {
+ return STOPPED
+ }
+ if ctx.exitErr != nil || len(ctx.processErrs) > 0 {
+ return FAILED
+ }
+ return FINISHED
+ default:
+ return ""
+ }
+}
+
+func formatPipelineResultError(exitErr error, processErrs []error) string {
+ if exitErr == nil && len(processErrs) == 0 {
+ return ""
+ }
+ return fmt.Sprintf("exit: %v, process: %v", exitErr, processErrs)
+}
+
func (ctx *Context) pushPipelineLog() {
if global.Env().IsDebug {
log.Info("received pipeline state change, id: ", ctx.Config.Name, ", state: ", ctx.runningState)
@@ -407,8 +445,8 @@ func (ctx *Context) pushPipelineLog() {
result := util.MapStr{
"success": ctx.exitErr == nil,
}
- if ctx.exitErr != nil || len(ctx.processErrs) > 0 {
- result["error"] = fmt.Sprintf("exit: %v, process: %v", ctx.exitErr, ctx.processErrs)
+ if errMsg := formatPipelineResultError(ctx.exitErr, ctx.processErrs); errMsg != "" {
+ result["error"] = errMsg
}
payload["result"] = result
}
@@ -418,5 +456,7 @@ func (ctx *Context) pushPipelineLog() {
},
}
- event.SaveLog(&eventData)
+ if err := event.SaveLog(&eventData); err != nil {
+ log.Errorf("failed to save pipeline log event, pipeline: %s, context: %s, err: %v", ctx.Config.Name, ctx.id, err)
+ }
}
diff --git a/core/pipeline/context_result_test.go b/core/pipeline/context_result_test.go
new file mode 100644
index 000000000..addb8d4b5
--- /dev/null
+++ b/core/pipeline/context_result_test.go
@@ -0,0 +1,59 @@
+package pipeline
+
+import (
+ "errors"
+ "testing"
+)
+
+func TestGetResultStateReturnsFinishedAfterStoppedCompletedRun(t *testing.T) {
+ ctx := AcquireContext(PipelineConfigV2{})
+ ctx.Started()
+ ctx.Finished()
+ ctx.Stopped()
+
+ if got := ctx.GetResultState(); got != FINISHED {
+ t.Fatalf("expected FINISHED result state, got %q", got)
+ }
+ if got := ctx.GetResultError(); got != "" {
+ t.Fatalf("expected empty result error, got %q", got)
+ }
+}
+
+func TestGetResultStateReturnsFailedAfterStoppedFailedRun(t *testing.T) {
+ ctx := AcquireContext(PipelineConfigV2{})
+ ctx.Started()
+ ctx.Failed(errors.New("boom"))
+ ctx.Stopped()
+
+ if got := ctx.GetResultState(); got != FAILED {
+ t.Fatalf("expected FAILED result state, got %q", got)
+ }
+ if got := ctx.GetResultError(); got == "" {
+ t.Fatal("expected result error for failed run")
+ }
+}
+
+func TestGetResultStateReturnsStoppedForManualStop(t *testing.T) {
+ ctx := AcquireContext(PipelineConfigV2{})
+ ctx.Started()
+ ctx.Stopping()
+ ctx.Stopped()
+
+ if got := ctx.GetResultState(); got != STOPPED {
+ t.Fatalf("expected STOPPED result state, got %q", got)
+ }
+ if got := ctx.GetResultError(); got != "" {
+ t.Fatalf("expected empty result error, got %q", got)
+ }
+}
+
+func TestGetResultErrorIncludesProcessErrors(t *testing.T) {
+ ctx := AcquireContext(PipelineConfigV2{})
+ ctx.Started()
+ ctx.RecordError(errors.New("slice failed"))
+ ctx.Finished()
+
+ if got := ctx.GetResultError(); got == "" {
+ t.Fatal("expected process error to be surfaced")
+ }
+}
diff --git a/core/queue/api.go b/core/queue/api.go
index e35182c8c..3740c045f 100755
--- a/core/queue/api.go
+++ b/core/queue/api.go
@@ -142,39 +142,49 @@ func AcquireConsumer(k *QueueConfig, consumer *ConsumerConfig, clientID string)
panic(errors.New("clientID can't be nil"))
}
- //check if the consumer is in fighting list
- if v, ok := consumersInFighting.Load(k.ID + consumer.Key()); ok {
- if v != clientID {
- //check the last touch time
+ fightingKey := k.ID + consumer.Key()
+
+ for {
+ reserved := false
+ currentOwner, loaded := consumersInFighting.LoadOrStore(fightingKey, clientID)
+ if loaded && currentOwner != clientID {
if consumer.ConsumeTimeoutInSeconds > 0 {
t := consumer.GetLastActiveTime()
if t != nil && int(time.Since(*t).Seconds()) > consumer.ConsumeTimeoutInSeconds {
- consumersInFighting.Delete(k.ID + consumer.Key())
- stats.Increment("consumer", k.ID, consumer.GetID(), "expired")
- //the consumer is in fighting and is already timeout
- return nil, errors.Errorf("consumer:%v is already in fighting list, but expired in: %v, remove it from the fighting list", consumer.Key(), time.Since(*t).Seconds())
+ if consumersInFighting.CompareAndDelete(fightingKey, currentOwner) {
+ stats.Increment("consumer", k.ID, consumer.GetID(), "expired")
+ }
+ continue
}
}
stats.Increment("consumer", k.ID, consumer.GetID(), "contend")
- //the consumer is in fighting list and the clientID is not the same
return nil, errors.New("the consumer is in fighting list")
}
- }
-
- handler := getAdvancedHandler(k)
- if handler != nil {
- v1, err := handler.AcquireConsumer(k, consumer)
- if err != nil {
- stats.Increment("consumer", k.ID, consumer.GetID(), "error_on_acquire")
- return nil, err
+ if !loaded {
+ reserved = true
}
- //add the consumer to the fighting list
- consumersInFighting.Store(k.ID+consumer.Key(), clientID)
- stats.Increment("consumer", k.ID, consumer.GetID(), "acquired")
- return v1, nil
+ handler := getAdvancedHandler(k)
+ if handler != nil {
+ acquired := false
+ defer func() {
+ if !acquired && reserved {
+ consumersInFighting.CompareAndDelete(fightingKey, clientID)
+ }
+ }()
+
+ v1, err := handler.AcquireConsumer(k, consumer)
+ if err != nil {
+ stats.Increment("consumer", k.ID, consumer.GetID(), "error_on_acquire")
+ return nil, err
+ }
+
+ acquired = true
+ stats.Increment("consumer", k.ID, consumer.GetID(), "acquired")
+ return v1, nil
+ }
+ panic(errors.New("handler is not registered"))
}
- panic(errors.New("handler is not registered"))
}
func ReleaseConsumer(k *QueueConfig, c *ConsumerConfig, consumer ConsumerAPI) error {
diff --git a/core/queue/api_test.go b/core/queue/api_test.go
new file mode 100644
index 000000000..829a2834a
--- /dev/null
+++ b/core/queue/api_test.go
@@ -0,0 +1,228 @@
+package queue
+
+import (
+ "errors"
+ "infini.sh/framework/core/stats"
+ "sync"
+ "testing"
+ "time"
+)
+
+type acquireConsumerTestHandler struct {
+ acquireFunc func(k *QueueConfig, consumer *ConsumerConfig) (ConsumerAPI, error)
+}
+
+func (h *acquireConsumerTestHandler) Name() string { return "test" }
+func (h *acquireConsumerTestHandler) Init(string) error { return nil }
+func (h *acquireConsumerTestHandler) Close(string) error { return nil }
+func (h *acquireConsumerTestHandler) GetStorageSize(string) uint64 { return 0 }
+func (h *acquireConsumerTestHandler) Destroy(string) error { return nil }
+func (h *acquireConsumerTestHandler) GetQueues() []string { return nil }
+func (h *acquireConsumerTestHandler) Push(string, []byte) error { return nil }
+func (h *acquireConsumerTestHandler) LatestOffset(*QueueConfig) Offset { return Offset{} }
+func (h *acquireConsumerTestHandler) GetOffset(*QueueConfig, *ConsumerConfig) (Offset, error) {
+ return Offset{}, nil
+}
+func (h *acquireConsumerTestHandler) DeleteOffset(*QueueConfig, *ConsumerConfig) error { return nil }
+func (h *acquireConsumerTestHandler) CommitOffset(*QueueConfig, *ConsumerConfig, Offset) (bool, error) {
+ return true, nil
+}
+func (h *acquireConsumerTestHandler) AcquireConsumer(k *QueueConfig, consumer *ConsumerConfig) (ConsumerAPI, error) {
+ if h.acquireFunc != nil {
+ return h.acquireFunc(k, consumer)
+ }
+ return &acquireConsumerTestConsumer{}, nil
+}
+func (h *acquireConsumerTestHandler) ReleaseConsumer(*QueueConfig, *ConsumerConfig, ConsumerAPI) error {
+ return nil
+}
+func (h *acquireConsumerTestHandler) AcquireProducer(*QueueConfig) (ProducerAPI, error) {
+ return nil, nil
+}
+func (h *acquireConsumerTestHandler) ReleaseProducer(*QueueConfig, ProducerAPI) error { return nil }
+
+type acquireConsumerTestConsumer struct{}
+
+func (c *acquireConsumerTestConsumer) Close() error { return nil }
+func (c *acquireConsumerTestConsumer) ResetOffset(int64, int64) error { return nil }
+func (c *acquireConsumerTestConsumer) FetchMessages(*Context, int) ([]Message, bool, error) {
+ return nil, false, nil
+}
+func (c *acquireConsumerTestConsumer) CommitOffset(Offset) error { return nil }
+
+type acquireConsumerTestStats struct {
+ mu sync.Mutex
+ timestamps map[string]time.Time
+}
+
+func (s *acquireConsumerTestStats) Increment(string, string) {}
+func (s *acquireConsumerTestStats) IncrementBy(string, string, int64) {}
+func (s *acquireConsumerTestStats) Decrement(string, string) {}
+func (s *acquireConsumerTestStats) DecrementBy(string, string, int64) {}
+func (s *acquireConsumerTestStats) Absolute(string, string, int64) {}
+func (s *acquireConsumerTestStats) Timing(string, string, int64) {}
+func (s *acquireConsumerTestStats) Gauge(string, string, int64) {}
+func (s *acquireConsumerTestStats) Stat(string, string) int64 { return 0 }
+func (s *acquireConsumerTestStats) StatsAll() string { return "" }
+func (s *acquireConsumerTestStats) RecordTimestamp(category, key string, value time.Time) {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ s.timestamps[category+"."+key] = value
+}
+func (s *acquireConsumerTestStats) GetTimestamp(category, key string) (time.Time, error) {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ v, ok := s.timestamps[category+"."+key]
+ if !ok {
+ return time.Time{}, errors.New("not found")
+ }
+ return v, nil
+}
+func (s *acquireConsumerTestStats) reset() {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ s.timestamps = map[string]time.Time{}
+}
+
+var acquireConsumerStatsOnce sync.Once
+var acquireConsumerStatsHandler = &acquireConsumerTestStats{timestamps: map[string]time.Time{}}
+
+func withTestQueueHandler(t *testing.T, handler AdvancedQueueAPI) {
+ t.Helper()
+ previousDefaultHandler := defaultHandler
+ previousConsumersInFighting := consumersInFighting
+ acquireConsumerStatsOnce.Do(func() {
+ stats.Register(acquireConsumerStatsHandler)
+ })
+ acquireConsumerStatsHandler.reset()
+ defaultHandler = handler
+ consumersInFighting = syncMapZero()
+ t.Cleanup(func() {
+ defaultHandler = previousDefaultHandler
+ consumersInFighting = previousConsumersInFighting
+ })
+}
+
+func syncMapZero() sync.Map {
+ return sync.Map{}
+}
+
+func TestAcquireConsumerStoresReservation(t *testing.T) {
+ withTestQueueHandler(t, &acquireConsumerTestHandler{})
+
+ q := &QueueConfig{ID: "queue-1", Name: "queue-1"}
+ c := &ConsumerConfig{Group: "group", Name: "consumer"}
+ c.ID = "consumer-1"
+
+ instance, err := AcquireConsumer(q, c, "client-1")
+ if err != nil {
+ t.Fatalf("expected acquire to succeed, got %v", err)
+ }
+ if instance == nil {
+ t.Fatal("expected consumer instance to be returned")
+ }
+ if owner, ok := consumersInFighting.Load(q.ID + c.Key()); !ok || owner != "client-1" {
+ t.Fatalf("expected fighting list reservation to be stored, got owner=%v exists=%v", owner, ok)
+ }
+}
+
+func TestAcquireConsumerRollsBackReservationOnError(t *testing.T) {
+ withTestQueueHandler(t, &acquireConsumerTestHandler{
+ acquireFunc: func(k *QueueConfig, consumer *ConsumerConfig) (ConsumerAPI, error) {
+ return nil, errors.New("boom")
+ },
+ })
+
+ q := &QueueConfig{ID: "queue-1", Name: "queue-1"}
+ c := &ConsumerConfig{Group: "group", Name: "consumer"}
+ c.ID = "consumer-1"
+
+ _, err := AcquireConsumer(q, c, "client-1")
+ if err == nil {
+ t.Fatal("expected acquire to fail")
+ }
+ if _, ok := consumersInFighting.Load(q.ID + c.Key()); ok {
+ t.Fatal("expected fighting list reservation to be rolled back")
+ }
+}
+
+func TestAcquireConsumerRollsBackReservationOnPanic(t *testing.T) {
+ withTestQueueHandler(t, &acquireConsumerTestHandler{
+ acquireFunc: func(k *QueueConfig, consumer *ConsumerConfig) (ConsumerAPI, error) {
+ panic("boom")
+ },
+ })
+
+ q := &QueueConfig{ID: "queue-1", Name: "queue-1"}
+ c := &ConsumerConfig{Group: "group", Name: "consumer"}
+ c.ID = "consumer-1"
+
+ defer func() {
+ if r := recover(); r == nil {
+ t.Fatal("expected acquire to panic")
+ }
+ if _, ok := consumersInFighting.Load(q.ID + c.Key()); ok {
+ t.Fatal("expected fighting list reservation to be rolled back after panic")
+ }
+ }()
+
+ _, _ = AcquireConsumer(q, c, "client-1")
+}
+
+func TestAcquireConsumerBlocksCompetingClient(t *testing.T) {
+ withTestQueueHandler(t, &acquireConsumerTestHandler{})
+
+ q := &QueueConfig{ID: "queue-1", Name: "queue-1"}
+ c := &ConsumerConfig{Group: "group", Name: "consumer"}
+ c.ID = "consumer-1"
+ consumersInFighting.Store(q.ID+c.Key(), "client-1")
+
+ _, err := AcquireConsumer(q, c, "client-2")
+ if err == nil || err.Error() != "the consumer is in fighting list" {
+ t.Fatalf("expected fighting list error, got %v", err)
+ }
+}
+
+func TestAcquireConsumerAllowsSameClientReentry(t *testing.T) {
+ withTestQueueHandler(t, &acquireConsumerTestHandler{})
+
+ q := &QueueConfig{ID: "queue-1", Name: "queue-1"}
+ c := &ConsumerConfig{Group: "group", Name: "consumer"}
+ c.ID = "consumer-1"
+ consumersInFighting.Store(q.ID+c.Key(), "client-1")
+
+ instance, err := AcquireConsumer(q, c, "client-1")
+ if err != nil {
+ t.Fatalf("expected same client to re-enter, got %v", err)
+ }
+ if instance == nil {
+ t.Fatal("expected consumer instance for same-client reentry")
+ }
+}
+
+func TestAcquireConsumerRetriesExpiredReservation(t *testing.T) {
+ withTestQueueHandler(t, &acquireConsumerTestHandler{})
+
+ q := &QueueConfig{ID: "queue-1", Name: "queue-1"}
+ c := &ConsumerConfig{
+ Group: "group",
+ Name: "consumer",
+ ConsumeTimeoutInSeconds: 1,
+ }
+ c.ID = "consumer-1"
+ c.KeepActive()
+ stale := time.Now().Add(-3 * time.Second)
+ stats.Timestamp("consumer", c.ID+".last_active", stale)
+ consumersInFighting.Store(q.ID+c.Key(), "client-2")
+
+ instance, err := AcquireConsumer(q, c, "client-1")
+ if err != nil {
+ t.Fatalf("expected expired reservation to be retried, got %v", err)
+ }
+ if instance == nil {
+ t.Fatal("expected consumer instance after expired reservation cleanup")
+ }
+ if owner, ok := consumersInFighting.Load(q.ID + c.Key()); !ok || owner != "client-1" {
+ t.Fatalf("expected ownership to move to client-1, got owner=%v exists=%v", owner, ok)
+ }
+}
diff --git a/core/queue/consumer_config.go b/core/queue/consumer_config.go
index 8572bcf47..2f0b3f328 100644
--- a/core/queue/consumer_config.go
+++ b/core/queue/consumer_config.go
@@ -172,7 +172,7 @@ func RemoveAllConsumers(qConfig *QueueConfig) (bool, error) {
log.Error(err)
return false, err
}
- log.Debugf("success delete all consumers for queue:%v", qConfig.ID)
+ log.Tracef("success delete all consumers for queue:%v", qConfig.ID)
return true, nil
}
diff --git a/core/queue/queue_config.go b/core/queue/queue_config.go
index 8d2671869..d03b2f64c 100644
--- a/core/queue/queue_config.go
+++ b/core/queue/queue_config.go
@@ -118,7 +118,7 @@ func RegisterConfig(cfg *QueueConfig) (preExists bool, err error) {
cfg.Created = time.Now().String()
- log.Debug("init new queue config:", cfg.ID, ",", cfg.Name)
+ log.Trace("init new queue config:", cfg.ID, ",", cfg.Name)
addCfgToCache(cfg)
diff --git a/core/security/password_challenge.go b/core/security/password_challenge.go
new file mode 100644
index 000000000..c007031cd
--- /dev/null
+++ b/core/security/password_challenge.go
@@ -0,0 +1,153 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package security
+
+import (
+ "errors"
+
+ "golang.org/x/crypto/bcrypt"
+ passwordchallenge "infini.sh/framework/core/security/passwordchallenge"
+ "infini.sh/framework/core/util"
+)
+
+const (
+ // PasswordChallengeMethod identifies the login flow returned by the challenge endpoint.
+ PasswordChallengeMethod = passwordchallenge.Method
+ // PasswordChallengeAlgorithm describes the verifier/proof derivation algorithm for clients.
+ PasswordChallengeAlgorithm = passwordchallenge.Algorithm
+ // PasswordChallengeIterations tells clients which PBKDF2 work factor to use.
+ PasswordChallengeIterations = passwordchallenge.Iterations
+)
+
+// LoginChallenge re-exports the framework challenge payload used by native account login.
+type LoginChallenge = passwordchallenge.Challenge
+
+// PasswordMaterial bundles the fields that apps need to persist after accepting a password.
+type PasswordMaterial struct {
+ Hash string
+ Salt string
+ Verifier string
+}
+
+// CanUsePasswordChallenge reports whether a native account already has challenge credentials.
+func CanUsePasswordChallenge(user *UserAccount) bool {
+ return user != nil && user.PasswordSalt != "" && user.PasswordVerifier != ""
+}
+
+// GeneratePasswordMaterial derives the bcrypt hash and challenge verifier fields for a password.
+func GeneratePasswordMaterial(password string) (*PasswordMaterial, error) {
+ hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
+ if err != nil {
+ return nil, err
+ }
+
+ // Store both the bcrypt hash for existing password checks and the derived
+ // verifier for challenge login so the two login modes stay in sync.
+ salt := util.GenerateSecureString(32)
+ verifier, err := DerivePasswordVerifier(password, salt)
+ if err != nil {
+ return nil, err
+ }
+
+ return &PasswordMaterial{
+ Hash: string(hash),
+ Salt: salt,
+ Verifier: verifier,
+ }, nil
+}
+
+// SetPassword updates both the legacy bcrypt hash and the challenge verifier material.
+func SetPassword(user *UserAccount, password string) error {
+ if user == nil {
+ return errors.New("user is nil")
+ }
+
+ material, err := GeneratePasswordMaterial(password)
+ if err != nil {
+ return err
+ }
+
+ user.Password = material.Hash
+ user.PasswordSalt = material.Salt
+ user.PasswordVerifier = material.Verifier
+ return nil
+}
+
+// EnsurePasswordChallenge derives challenge material for older accounts without changing the bcrypt hash.
+func EnsurePasswordChallenge(user *UserAccount, password string) error {
+ if user == nil {
+ return errors.New("user is nil")
+ }
+ if CanUsePasswordChallenge(user) {
+ return nil
+ }
+
+ // This is used as an in-place upgrade path for older native accounts that only
+ // have a bcrypt password hash from before challenge login was introduced.
+ salt := util.GenerateSecureString(32)
+ verifier, err := DerivePasswordVerifier(password, salt)
+ if err != nil {
+ return err
+ }
+
+ user.PasswordSalt = salt
+ user.PasswordVerifier = verifier
+ return nil
+}
+
+// VerifyPassword validates the plain password against the stored bcrypt hash.
+func VerifyPassword(user *UserAccount, password string) error {
+ if user == nil {
+ return errors.New("user is nil")
+ }
+ if user.Password == "" {
+ return errors.New("password is not set")
+ }
+ return bcrypt.CompareHashAndPassword([]byte(user.Password), []byte(password))
+}
+
+// DerivePasswordVerifier converts a password and salt into the stored challenge verifier.
+func DerivePasswordVerifier(password, salt string) (string, error) {
+ return passwordchallenge.DeriveVerifier(password, salt)
+}
+
+// BuildPasswordProof creates the challenge response that clients send to /account/login.
+func BuildPasswordProof(verifier, subject, challengeID, nonce string) (string, error) {
+ return passwordchallenge.BuildProof(verifier, subject, challengeID, nonce)
+}
+
+// VerifyPasswordProof checks whether a submitted proof matches the stored verifier.
+func VerifyPasswordProof(verifier, subject, challengeID, nonce, proof string) bool {
+ return passwordchallenge.VerifyProof(verifier, subject, challengeID, nonce, proof)
+}
+
+// NewLoginChallenge allocates a one-time challenge bound to the requested login subject.
+func NewLoginChallenge(subject string) LoginChallenge {
+ return passwordchallenge.New(subject)
+}
+
+// ConsumeLoginChallenge validates and removes a one-time challenge after it is used.
+func ConsumeLoginChallenge(challengeID, subject string) (LoginChallenge, error) {
+ return passwordchallenge.Consume(challengeID, subject)
+}
diff --git a/core/security/password_challenge_test.go b/core/security/password_challenge_test.go
new file mode 100644
index 000000000..c2f68a145
--- /dev/null
+++ b/core/security/password_challenge_test.go
@@ -0,0 +1,108 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package security
+
+import "testing"
+
+// Setting a password should populate both the legacy bcrypt path and the new
+// challenge-login material so either login mode can succeed afterward.
+func TestSetPasswordPopulatesChallengeFields(t *testing.T) {
+ user := &UserAccount{}
+ if err := SetPassword(user, "StrongPassw0rd!"); err != nil {
+ t.Fatalf("set password: %v", err)
+ }
+
+ if user.Password == "" {
+ t.Fatal("expected password hash to be set")
+ }
+ if user.PasswordSalt == "" {
+ t.Fatal("expected password salt to be set")
+ }
+ if user.PasswordVerifier == "" {
+ t.Fatal("expected password verifier to be set")
+ }
+ if err := VerifyPassword(user, "StrongPassw0rd!"); err != nil {
+ t.Fatalf("verify password: %v", err)
+ }
+}
+
+// Existing challenge material should be stable when an account is already upgraded.
+func TestEnsurePasswordChallengePreservesExistingVerifier(t *testing.T) {
+ user := &UserAccount{}
+ if err := SetPassword(user, "StrongPassw0rd!"); err != nil {
+ t.Fatalf("set password: %v", err)
+ }
+
+ originalSalt := user.PasswordSalt
+ originalVerifier := user.PasswordVerifier
+ if err := EnsurePasswordChallenge(user, "AnotherStrongPassw0rd!"); err != nil {
+ t.Fatalf("ensure password challenge: %v", err)
+ }
+
+ if user.PasswordSalt != originalSalt {
+ t.Fatal("expected existing password salt to be preserved")
+ }
+ if user.PasswordVerifier != originalVerifier {
+ t.Fatal("expected existing password verifier to be preserved")
+ }
+}
+
+// The framework wrapper should produce proofs compatible with the lower-level package.
+func TestPasswordChallengeProofRoundTrip(t *testing.T) {
+ user := &UserAccount{}
+ login := "admin@example.org"
+ password := "StrongPassw0rd!"
+
+ if err := SetPassword(user, password); err != nil {
+ t.Fatalf("set password: %v", err)
+ }
+
+ challenge := NewLoginChallenge(login)
+ proof, err := BuildPasswordProof(user.PasswordVerifier, login, challenge.ID, challenge.Nonce)
+ if err != nil {
+ t.Fatalf("build password proof: %v", err)
+ }
+
+ if !VerifyPasswordProof(user.PasswordVerifier, login, challenge.ID, challenge.Nonce, proof) {
+ t.Fatal("expected challenge proof to validate")
+ }
+}
+
+// Legacy accounts that only had a bcrypt hash should become challenge-capable in place.
+func TestEnsurePasswordChallengePopulatesLegacyAccount(t *testing.T) {
+ user := &UserAccount{Password: "existing-bcrypt-hash"}
+ if err := EnsurePasswordChallenge(user, "StrongPassw0rd!"); err != nil {
+ t.Fatalf("ensure password challenge: %v", err)
+ }
+
+ if user.PasswordSalt == "" {
+ t.Fatal("expected password salt to be populated")
+ }
+ if user.PasswordVerifier == "" {
+ t.Fatal("expected password verifier to be populated")
+ }
+ if !CanUsePasswordChallenge(user) {
+ t.Fatal("expected legacy account to become challenge-capable")
+ }
+}
diff --git a/core/security/passwordchallenge/password_challenge.go b/core/security/passwordchallenge/password_challenge.go
new file mode 100644
index 000000000..f16190bdb
--- /dev/null
+++ b/core/security/passwordchallenge/password_challenge.go
@@ -0,0 +1,184 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package passwordchallenge
+
+import (
+ "crypto/hmac"
+ "crypto/sha256"
+ "encoding/hex"
+ "errors"
+ "strings"
+ "sync"
+ "time"
+
+ "golang.org/x/crypto/pbkdf2"
+ "infini.sh/framework/core/util"
+)
+
+const (
+ // Method identifies the password challenge login flow returned by the challenge endpoint.
+ Method = "challenge"
+ // Algorithm describes the verifier/proof derivation algorithm that clients must use.
+ Algorithm = "PBKDF2-SHA256"
+ // Iterations is the PBKDF2 work factor shared with clients during challenge negotiation.
+ Iterations = 120000
+ // keyLength is the derived key size used for both the stored verifier and request proof.
+ keyLength = 32
+ // DefaultTTL is the default lifetime of a login challenge before it must be re-issued.
+ DefaultTTL = 5 * time.Minute
+)
+
+// Challenge carries the one-time identifiers clients need to build a password proof locally.
+type Challenge struct {
+ ID string
+ // Subject keeps the challenge bound to the login identity it was issued for.
+ Subject string
+ // Nonce is the random per-challenge input mixed into the client proof.
+ Nonce string
+ // ExpireAt marks when the one-time challenge stops being valid.
+ ExpireAt time.Time
+}
+
+// StoreOptions configures the lifetime of issued login challenges.
+type StoreOptions struct {
+ // TTL overrides the default challenge lifetime for this store instance.
+ TTL time.Duration
+}
+
+// Store tracks outstanding login challenges until they are consumed or expire.
+type Store struct {
+ mu sync.Mutex
+ ttl time.Duration
+ challenges map[string]Challenge
+}
+
+var defaultStore = NewStore(StoreOptions{})
+
+// NewStore creates an in-memory challenge store with the requested TTL.
+func NewStore(options StoreOptions) *Store {
+ ttl := options.TTL
+ if ttl <= 0 {
+ ttl = DefaultTTL
+ }
+ return &Store{
+ ttl: ttl,
+ challenges: map[string]Challenge{},
+ }
+}
+
+// DeriveVerifier turns a password and salt into the verifier stored on the account record.
+func DeriveVerifier(password, salt string) (string, error) {
+ if password == "" {
+ return "", errors.New("password is empty")
+ }
+ if salt == "" {
+ return "", errors.New("password salt is empty")
+ }
+ key := pbkdf2.Key([]byte(password), []byte(salt), Iterations, keyLength, sha256.New)
+ return hex.EncodeToString(key), nil
+}
+
+// BuildProof derives the one-time challenge response that clients submit to /account/login.
+func BuildProof(verifier, subject, challengeID, nonce string) (string, error) {
+ key, err := hex.DecodeString(verifier)
+ if err != nil {
+ return "", err
+ }
+ mac := hmac.New(sha256.New, key)
+ mac.Write([]byte(subject))
+ mac.Write([]byte(":"))
+ mac.Write([]byte(challengeID))
+ mac.Write([]byte(":"))
+ mac.Write([]byte(nonce))
+ return hex.EncodeToString(mac.Sum(nil)), nil
+}
+
+// VerifyProof compares a submitted proof against the expected proof for this challenge tuple.
+func VerifyProof(verifier, subject, challengeID, nonce, proof string) bool {
+ expected, err := BuildProof(verifier, subject, challengeID, nonce)
+ if err != nil {
+ return false
+ }
+ expectedBytes, err := hex.DecodeString(expected)
+ if err != nil {
+ return false
+ }
+ proofBytes, err := hex.DecodeString(strings.ToLower(proof))
+ if err != nil {
+ return false
+ }
+ return hmac.Equal(expectedBytes, proofBytes)
+}
+
+// New issues a login challenge from the default store.
+func New(subject string) Challenge {
+ return defaultStore.New(subject)
+}
+
+// Consume loads and invalidates a login challenge from the default store.
+func Consume(challengeID, subject string) (Challenge, error) {
+ return defaultStore.Consume(challengeID, subject)
+}
+
+// New allocates a fresh challenge for the provided subject.
+func (store *Store) New(subject string) Challenge {
+ now := time.Now()
+ store.mu.Lock()
+ defer store.mu.Unlock()
+
+ for id, challenge := range store.challenges {
+ if challenge.ExpireAt.Before(now) {
+ delete(store.challenges, id)
+ }
+ }
+
+ challenge := Challenge{
+ ID: util.GenerateSecureString(32),
+ Subject: subject,
+ Nonce: util.GenerateSecureString(32),
+ ExpireAt: now.Add(store.ttl),
+ }
+ store.challenges[challenge.ID] = challenge
+ return challenge
+}
+
+// Consume validates the subject and TTL, then invalidates the one-time challenge.
+func (store *Store) Consume(challengeID, subject string) (Challenge, error) {
+ store.mu.Lock()
+ defer store.mu.Unlock()
+
+ challenge, ok := store.challenges[challengeID]
+ if !ok {
+ return Challenge{}, errors.New("login challenge is invalid")
+ }
+ delete(store.challenges, challengeID)
+
+ if challenge.ExpireAt.Before(time.Now()) {
+ return Challenge{}, errors.New("login challenge expired")
+ }
+ if challenge.Subject != subject {
+ return Challenge{}, errors.New("login challenge does not match user")
+ }
+ return challenge, nil
+}
diff --git a/core/security/passwordchallenge/password_challenge_test.go b/core/security/passwordchallenge/password_challenge_test.go
new file mode 100644
index 000000000..8b0ce6aac
--- /dev/null
+++ b/core/security/passwordchallenge/password_challenge_test.go
@@ -0,0 +1,79 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package passwordchallenge
+
+import (
+ "testing"
+ "time"
+)
+
+// Proof generation and verification need to round-trip because this package defines the
+// wire contract shared between the framework login endpoint and upgraded clients.
+func TestPasswordChallengeProofRoundTrip(t *testing.T) {
+ verifier, err := DeriveVerifier("admin", "salt-123")
+ if err != nil {
+ t.Fatalf("derive verifier: %v", err)
+ }
+
+ challenge := New("admin")
+ proof, err := BuildProof(verifier, "admin", challenge.ID, challenge.Nonce)
+ if err != nil {
+ t.Fatalf("build proof: %v", err)
+ }
+
+ if !VerifyProof(verifier, "admin", challenge.ID, challenge.Nonce, proof) {
+ t.Fatal("expected password proof to validate")
+ }
+}
+
+// Challenges are bound to the requested login subject and must not be replayed for others.
+func TestConsumeRejectsWrongSubject(t *testing.T) {
+ store := NewStore(StoreOptions{})
+ challenge := store.New("admin")
+
+ if _, err := store.Consume(challenge.ID, "guest"); err == nil {
+ t.Fatal("expected challenge subject mismatch to fail")
+ }
+}
+
+// Empty input should be rejected up front to avoid persisting or comparing invalid verifiers.
+func TestDeriveVerifierRejectsEmptyInput(t *testing.T) {
+ if _, err := DeriveVerifier("", "salt-123"); err == nil {
+ t.Fatal("expected empty password to fail")
+ }
+ if _, err := DeriveVerifier("admin", ""); err == nil {
+ t.Fatal("expected empty salt to fail")
+ }
+}
+
+// Expiration keeps the one-time challenge store bounded and prevents stale proof reuse.
+func TestConsumeRejectsExpiredChallenge(t *testing.T) {
+ store := NewStore(StoreOptions{TTL: time.Millisecond})
+ challenge := store.New("admin")
+
+ time.Sleep(5 * time.Millisecond)
+ if _, err := store.Consume(challenge.ID, "admin"); err == nil {
+ t.Fatal("expected expired challenge to fail")
+ }
+}
diff --git a/core/security/replay/replay.go b/core/security/replay/replay.go
new file mode 100644
index 000000000..b5a630eaf
--- /dev/null
+++ b/core/security/replay/replay.go
@@ -0,0 +1,213 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package replay
+
+import (
+ "crypto/sha256"
+ "encoding/hex"
+ "fmt"
+ "net/http"
+ pathutil "path"
+ "strings"
+ "sync"
+ "time"
+
+ "infini.sh/framework/core/util"
+)
+
+const (
+ // HeaderName is the HTTP header clients use to submit a one-time replay nonce.
+ HeaderName = "X-Request-Nonce"
+ // DefaultTTL is the default lifetime of an issued replay nonce.
+ DefaultTTL = 30 * time.Second
+)
+
+// SubjectExtractor derives the caller identity that a replay nonce should be bound to.
+type SubjectExtractor func(r *http.Request) string
+
+// StoreOptions configures replay nonce retention and subject binding behavior.
+type StoreOptions struct {
+ TTL time.Duration
+ SubjectExtractor SubjectExtractor
+}
+
+type nonceRecord struct {
+ Subject string
+ Method string
+ Path string
+ ExpiresAt time.Time
+}
+
+// Store tracks issued replay nonces until they are consumed or expire.
+type Store struct {
+ mu sync.Mutex
+ ttl time.Duration
+ subjectExtractor SubjectExtractor
+ records map[string]nonceRecord
+}
+
+var defaultStore = NewStore(StoreOptions{})
+
+// NewStore creates an in-memory replay store with optional TTL and subject extraction overrides.
+func NewStore(options StoreOptions) *Store {
+ ttl := options.TTL
+ if ttl <= 0 {
+ ttl = DefaultTTL
+ }
+ extractor := options.SubjectExtractor
+ if extractor == nil {
+ extractor = DefaultSubjectExtractor
+ }
+ return &Store{
+ ttl: ttl,
+ subjectExtractor: extractor,
+ records: map[string]nonceRecord{},
+ }
+}
+
+// IssueReplayNonce issues a nonce from the default store for the requested method/path scope.
+func IssueReplayNonce(r *http.Request, method, requestPath string) (string, time.Duration, error) {
+ return defaultStore.IssueReplayNonce(r, method, requestPath)
+}
+
+// ValidateAndConsumeReplayNonce validates a nonce from the default store and deletes it on success.
+func ValidateAndConsumeReplayNonce(r *http.Request) error {
+ return defaultStore.ValidateAndConsumeReplayNonce(r)
+}
+
+// DefaultSubjectExtractor binds anonymous callers together and authenticated callers to their
+// Authorization header so replay nonces cannot be replayed across credential contexts.
+func DefaultSubjectExtractor(r *http.Request) string {
+ if r == nil {
+ return "anonymous"
+ }
+ authorizationHeader := strings.TrimSpace(r.Header.Get("Authorization"))
+ if authorizationHeader == "" {
+ return "anonymous"
+ }
+ // Bind the nonce to the caller's authorization material so a replay token issued for one
+ // authenticated context cannot be reused with a different credential set.
+ sum := sha256.Sum256([]byte(authorizationHeader))
+ return hex.EncodeToString(sum[:])
+}
+
+// IssueReplayNonce stores a nonce that is scoped to the caller, HTTP method, and request path.
+func (store *Store) IssueReplayNonce(r *http.Request, method, requestPath string) (string, time.Duration, error) {
+ normalizedMethod, normalizedPath, err := normalizeScope(method, requestPath)
+ if err != nil {
+ return "", 0, err
+ }
+
+ nonce := util.GenerateSecureString(32)
+ if nonce == "" {
+ return "", 0, fmt.Errorf("failed to generate replay nonce")
+ }
+
+ subject := store.extractSubject(r)
+ expiresAt := time.Now().Add(store.ttl)
+
+ store.mu.Lock()
+ defer store.mu.Unlock()
+ store.cleanupExpiredLocked(time.Now())
+ store.records[nonce] = nonceRecord{
+ Subject: subject,
+ Method: normalizedMethod,
+ Path: normalizedPath,
+ ExpiresAt: expiresAt,
+ }
+ return nonce, store.ttl, nil
+}
+
+// ValidateAndConsumeReplayNonce accepts a nonce only once and only for the original request scope.
+func (store *Store) ValidateAndConsumeReplayNonce(r *http.Request) error {
+ if r == nil {
+ return fmt.Errorf("request can not be nil")
+ }
+
+ nonce := strings.TrimSpace(r.Header.Get(HeaderName))
+ if nonce == "" {
+ return fmt.Errorf("missing replay nonce")
+ }
+
+ subject := store.extractSubject(r)
+ method, requestPath, err := normalizeScope(r.Method, r.URL.Path)
+ if err != nil {
+ return err
+ }
+
+ now := time.Now()
+ store.mu.Lock()
+ defer store.mu.Unlock()
+ store.cleanupExpiredLocked(now)
+
+ record, ok := store.records[nonce]
+ if !ok {
+ return fmt.Errorf("replay nonce is invalid or expired")
+ }
+ delete(store.records, nonce)
+
+ if record.Subject != subject || record.Method != method || record.Path != requestPath {
+ return fmt.Errorf("replay nonce does not match request context")
+ }
+
+ return nil
+}
+
+func (store *Store) extractSubject(r *http.Request) string {
+ if store == nil || store.subjectExtractor == nil {
+ return DefaultSubjectExtractor(r)
+ }
+ return store.subjectExtractor(r)
+}
+
+func (store *Store) cleanupExpiredLocked(now time.Time) {
+ for nonce, record := range store.records {
+ if now.After(record.ExpiresAt) {
+ delete(store.records, nonce)
+ }
+ }
+}
+
+func normalizeScope(method, requestPath string) (string, string, error) {
+ normalizedMethod := strings.ToUpper(strings.TrimSpace(method))
+ switch normalizedMethod {
+ case http.MethodPost, http.MethodPut, http.MethodDelete:
+ default:
+ return "", "", fmt.Errorf("unsupported replay-protected method [%s]", method)
+ }
+
+ normalizedPath := strings.TrimSpace(requestPath)
+ if normalizedPath == "" {
+ return "", "", fmt.Errorf("request path can not be empty")
+ }
+ if !strings.HasPrefix(normalizedPath, "/") {
+ normalizedPath = "/" + normalizedPath
+ }
+ normalizedPath = pathutil.Clean(normalizedPath)
+ if normalizedPath == "." {
+ normalizedPath = "/"
+ }
+
+ return normalizedMethod, normalizedPath, nil
+}
diff --git a/core/security/replay/replay_test.go b/core/security/replay/replay_test.go
new file mode 100644
index 000000000..414cd79b5
--- /dev/null
+++ b/core/security/replay/replay_test.go
@@ -0,0 +1,128 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package replay
+
+import (
+ "net/http"
+ "net/http/httptest"
+ "testing"
+ "time"
+)
+
+// A nonce is one-time use by design, so any second validation attempt must fail.
+func TestReplayNonceCanOnlyBeUsedOnce(t *testing.T) {
+ store := NewStore(StoreOptions{})
+ req := httptest.NewRequest(http.MethodPost, "https://console.local/account/login", nil)
+
+ nonce, _, err := store.IssueReplayNonce(req, http.MethodPost, "/account/login")
+ if err != nil {
+ t.Fatalf("issue replay nonce failed: %v", err)
+ }
+
+ req.Header.Set(HeaderName, nonce)
+ if err := store.ValidateAndConsumeReplayNonce(req); err != nil {
+ t.Fatalf("expected first nonce use to succeed: %v", err)
+ }
+ if err := store.ValidateAndConsumeReplayNonce(req); err == nil {
+ t.Fatal("expected second nonce use to be rejected")
+ }
+}
+
+// Replay tokens should stay bound to the caller identity, not just the raw path/method tuple.
+func TestReplayNonceBindsToAuthorizationHeader(t *testing.T) {
+ store := NewStore(StoreOptions{})
+ issueReq := httptest.NewRequest(http.MethodPut, "https://console.local/credential/test", nil)
+ issueReq.Header.Set("Authorization", "Bearer token-a")
+
+ nonce, _, err := store.IssueReplayNonce(issueReq, http.MethodPut, "/credential/test")
+ if err != nil {
+ t.Fatalf("issue replay nonce failed: %v", err)
+ }
+
+ useReq := httptest.NewRequest(http.MethodPut, "https://console.local/credential/test", nil)
+ useReq.Header.Set(HeaderName, nonce)
+ useReq.Header.Set("Authorization", "Bearer token-b")
+ if err := store.ValidateAndConsumeReplayNonce(useReq); err == nil {
+ t.Fatal("expected nonce bound to a different authorization header to fail")
+ }
+}
+
+// The scope includes HTTP method so a nonce issued for one mutation cannot authorize another.
+func TestReplayNonceBindsToPathAndMethod(t *testing.T) {
+ store := NewStore(StoreOptions{})
+ issueReq := httptest.NewRequest(http.MethodPost, "https://console.local/setup/_initialize", nil)
+
+ nonce, _, err := store.IssueReplayNonce(issueReq, http.MethodPost, "/setup/_initialize")
+ if err != nil {
+ t.Fatalf("issue replay nonce failed: %v", err)
+ }
+
+ useReq := httptest.NewRequest(http.MethodPut, "https://console.local/setup/_initialize", nil)
+ useReq.Header.Set(HeaderName, nonce)
+ if err := store.ValidateAndConsumeReplayNonce(useReq); err == nil {
+ t.Fatal("expected nonce with mismatched method to fail")
+ }
+}
+
+// Anonymous callers still need a stable default subject so unauthenticated setup flows work.
+func TestDefaultSubjectExtractorFallsBackToAnonymous(t *testing.T) {
+ req := httptest.NewRequest(http.MethodPost, "https://console.local/account/login", nil)
+ if got := DefaultSubjectExtractor(req); got != "anonymous" {
+ t.Fatalf("expected anonymous subject, got %q", got)
+ }
+}
+
+// Path normalization lets clients request a nonce with equivalent path forms safely.
+func TestReplayNonceNormalizesPath(t *testing.T) {
+ store := NewStore(StoreOptions{})
+ issueReq := httptest.NewRequest(http.MethodPost, "https://console.local/account/login", nil)
+
+ nonce, _, err := store.IssueReplayNonce(issueReq, http.MethodPost, "account/../account/login")
+ if err != nil {
+ t.Fatalf("issue replay nonce failed: %v", err)
+ }
+
+ useReq := httptest.NewRequest(http.MethodPost, "https://console.local/account/login", nil)
+ useReq.Header.Set(HeaderName, nonce)
+ if err := store.ValidateAndConsumeReplayNonce(useReq); err != nil {
+ t.Fatalf("expected normalized path to validate: %v", err)
+ }
+}
+
+// Expired nonces should be rejected even if the caller, method, and path still match.
+func TestReplayNonceExpires(t *testing.T) {
+ store := NewStore(StoreOptions{TTL: time.Millisecond})
+ req := httptest.NewRequest(http.MethodPost, "https://console.local/account/login", nil)
+
+ nonce, _, err := store.IssueReplayNonce(req, http.MethodPost, "/account/login")
+ if err != nil {
+ t.Fatalf("issue replay nonce failed: %v", err)
+ }
+
+ time.Sleep(5 * time.Millisecond)
+ req.Header.Set(HeaderName, nonce)
+ if err := store.ValidateAndConsumeReplayNonce(req); err == nil {
+ t.Fatal("expected expired nonce to be rejected")
+ }
+}
diff --git a/core/security/role_registry.go b/core/security/role_registry.go
index 18ec757db..f9b0fb724 100644
--- a/core/security/role_registry.go
+++ b/core/security/role_registry.go
@@ -160,8 +160,8 @@ func (rr *RoleRegistry) GetPermissionsForRole(role string) ([]PermissionKey, boo
}
func GetAllPermissionsForUser(user *UserSessionInfo) []PermissionKey {
- if user==nil{
- return []PermissionKey{}
+ if user == nil {
+ return []PermissionKey{}
}
permissions := user.GetPermissionKeys()
@@ -194,8 +194,8 @@ func GetAllPermissionsForUser(user *UserSessionInfo) []PermissionKey {
}
func getPermissionKeysByUser(user *UserSessionInfo) ([]PermissionKey, error) {
- if user==nil{
- return []PermissionKey{},nil
+ if user == nil {
+ return []PermissionKey{}, nil
}
ctx1 := context.Background()
diff --git a/core/security/service_registry.go b/core/security/service_registry.go
index 5173cebb1..bc8b6595c 100644
--- a/core/security/service_registry.go
+++ b/core/security/service_registry.go
@@ -21,6 +21,12 @@ type AuthorizationBackend interface {
GetPermissionKeysByRoles(ctx context.Context, roles []string) []PermissionKey
}
+// AccountPasswordLoginProvider lets applications keep their own password-auth realms
+// while reusing the shared framework account login handler and session issuance.
+type AccountPasswordLoginProvider interface {
+ AuthenticateByPassword(login, password string) (*UserSessionInfo, error)
+}
+
var authorizationBackendProviders = sync.Map{}
func RegisterAuthorizationProvider(name string, provider AuthorizationBackend) {
@@ -33,6 +39,12 @@ func RegisterAuthenticationProvider(name string, provider AuthenticationBackend)
authenticationBackendBackendProviders.Store(name, provider)
}
+var accountPasswordLoginProviders = sync.Map{}
+
+func RegisterAccountPasswordLoginProvider(name string, provider AccountPasswordLoginProvider) {
+ accountPasswordLoginProviders.Store(name, provider)
+}
+
func MustGetAuthenticationProvider(provider string) AuthenticationBackend {
value, ok := authenticationBackendBackendProviders.Load(provider)
if ok {
@@ -97,5 +109,35 @@ func GetUserByLogin(login string) (bool, *UserAccount, error) {
return false, nil, errors.New("no AuthenticationBackend was found")
}
- return false, nil, errors.New("not found")
+ return false, nil, nil
+}
+
+// AuthenticateAccountPasswordLogin tries application-provided password login providers
+// after the native framework account path has either not matched or not succeeded.
+func AuthenticateAccountPasswordLogin(login, password string) (*UserSessionInfo, error) {
+ var out *UserSessionInfo
+ var lastErr error
+
+ accountPasswordLoginProviders.Range(func(key, value any) bool {
+ provider, ok := value.(AccountPasswordLoginProvider)
+ if !ok {
+ return true
+ }
+
+ sessionUser, err := provider.AuthenticateByPassword(login, password)
+ if err != nil {
+ lastErr = err
+ return true
+ }
+ if sessionUser != nil {
+ out = sessionUser
+ return false
+ }
+ return true
+ })
+
+ if out != nil {
+ return out, nil
+ }
+ return nil, lastErr
}
diff --git a/core/security/session.go b/core/security/session.go
index cf455adf2..f3a21a83d 100644
--- a/core/security/session.go
+++ b/core/security/session.go
@@ -7,6 +7,7 @@ package security
import (
"fmt"
"net/http"
+ "sync"
"time"
"github.com/golang-jwt/jwt/v4"
@@ -16,11 +17,22 @@ import (
)
const UserAccessTokenSessionName = "user_session_access_token"
+const UserAccessTokenTTL = 24 * time.Hour
+
+// SessionTokenResponseDecorator lets applications enrich the shared login/refresh
+// response with app-specific fields while reusing the framework session pipeline.
+type SessionTokenResponseDecorator func(token map[string]interface{}, user *UserSessionInfo)
+
+var sessionTokenResponseDecorators = sync.Map{}
func init() {
RegisterHTTPAuthFilterProviderWithPriority("session_token", byAccessTokenSession, 10)
}
+func RegisterSessionTokenResponseDecorator(name string, decorator SessionTokenResponseDecorator) {
+ sessionTokenResponseDecorators.Store(name, decorator)
+}
+
func byAccessTokenSession(w http.ResponseWriter, r *http.Request) (claims *UserClaims, err error) {
exists, sessToken := api.GetSession(w, r, UserAccessTokenSessionName)
if !exists || sessToken == nil {
@@ -65,7 +77,7 @@ func byAccessTokenSession(w http.ResponseWriter, r *http.Request) (claims *UserC
func AddUserToSession(w http.ResponseWriter, r *http.Request, user *UserSessionInfo) (error, map[string]interface{}) {
if user == nil {
- panic("invalid user")
+ return errors.NewWithHTTPCode(http.StatusUnauthorized, "invalid user"), nil
}
// Generate access token
@@ -89,7 +101,7 @@ func GenerateJWTAccessToken(user *UserSessionInfo) (map[string]interface{}, erro
token1 := jwt.NewWithClaims(jwt.SigningMethodHS256, UserClaims{
UserSessionInfo: user,
RegisteredClaims: &jwt.RegisteredClaims{
- ExpiresAt: jwt.NewNumericDate(time.Now().Add(24 * time.Hour)),
+ ExpiresAt: jwt.NewNumericDate(time.Now().Add(UserAccessTokenTTL)),
},
})
@@ -105,7 +117,7 @@ func GenerateJWTAccessToken(user *UserSessionInfo) (map[string]interface{}, erro
data = util.MapStr{
"access_token": tokenString,
- "expire_in": time.Now().Unix() + 86400, //24h
+ "expire_in": time.Now().Unix() + int64(UserAccessTokenTTL/time.Second),
}
data["status"] = "ok"
@@ -113,3 +125,54 @@ func GenerateJWTAccessToken(user *UserSessionInfo) (map[string]interface{}, erro
return data, err
}
+
+// DecorateSessionTokenResponse keeps framework-issued account responses directly
+// consumable by existing console clients while auth flows converge on framework.
+func DecorateSessionTokenResponse(token map[string]interface{}, user *UserSessionInfo) {
+ if token == nil || user == nil {
+ return
+ }
+
+ if expiresAt := tokenExpiresAtUnix(token["expire_in"]); expiresAt > 0 {
+ token["expires_at"] = expiresAt
+
+ remaining := expiresAt - time.Now().Unix()
+ if remaining < 0 {
+ remaining = 0
+ }
+ token["expire_in"] = remaining
+ }
+
+ token["username"] = user.Login
+ token["id"] = user.UserID
+ token["roles"] = append([]string(nil), user.Roles...)
+ token["privilege"] = GetAllPermissionsForUser(user)
+ applySessionTokenResponseDecorators(token, user)
+}
+
+func tokenExpiresAtUnix(value interface{}) int64 {
+ switch v := value.(type) {
+ case int64:
+ return v
+ case int:
+ return int64(v)
+ case int32:
+ return int64(v)
+ case float64:
+ return int64(v)
+ case float32:
+ return int64(v)
+ default:
+ return 0
+ }
+}
+
+func applySessionTokenResponseDecorators(token map[string]interface{}, user *UserSessionInfo) {
+ sessionTokenResponseDecorators.Range(func(key, value any) bool {
+ decorator, ok := value.(SessionTokenResponseDecorator)
+ if ok {
+ decorator(token, user)
+ }
+ return true
+ })
+}
diff --git a/core/security/session_test.go b/core/security/session_test.go
new file mode 100644
index 000000000..aaa1d5cba
--- /dev/null
+++ b/core/security/session_test.go
@@ -0,0 +1,36 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package security
+
+import "testing"
+
+func TestAddUserToSessionRejectsNilUser(t *testing.T) {
+ err, token := AddUserToSession(nil, nil, nil)
+ if err == nil {
+ t.Fatal("expected nil user to be rejected")
+ }
+ if token != nil {
+ t.Fatalf("expected no token for nil user, got %+v", token)
+ }
+}
diff --git a/core/security/user_profile.go b/core/security/user_profile.go
index cd6936fe5..35155b0e2 100644
--- a/core/security/user_profile.go
+++ b/core/security/user_profile.go
@@ -33,10 +33,12 @@ type User struct {
type UserAccount struct {
orm.ORMObjectBase
- Name string `json:"name,omitempty" elastic_mapping:"name: { type: keyword }" validate:"required" `
- Email string `json:"email,omitempty" elastic_mapping:"email: { type: keyword }" validate:"required|email" ` //unique
- Roles []string `json:"roles,omitempty" elastic_mapping:"roles: { type: keyword }"`
- Password string `json:"password,omitempty" elastic_mapping:"password: { type: keyword }"`
+ Name string `json:"name,omitempty" elastic_mapping:"name: { type: keyword }" validate:"required" `
+ Email string `json:"email,omitempty" elastic_mapping:"email: { type: keyword }" validate:"required|email" ` //unique
+ Roles []string `json:"roles,omitempty" elastic_mapping:"roles: { type: keyword }"`
+ Password string `json:"password,omitempty" elastic_mapping:"password: { type: keyword }"` // Bcrypt hash used by the existing password-login flow.
+ PasswordSalt string `json:"password_salt,omitempty" elastic_mapping:"password_salt: { type: keyword }"` // Per-user salt exposed to clients during challenge login.
+ PasswordVerifier string `json:"password_verifier,omitempty" elastic_mapping:"password_verifier: { type: keyword }"` // Server-side verifier used to validate challenge proofs.
}
type UserProfile struct {
diff --git a/core/security/user_session.go b/core/security/user_session.go
index 3ec473b82..faef11760 100644
--- a/core/security/user_session.go
+++ b/core/security/user_session.go
@@ -56,7 +56,7 @@ type UserSessionInfo struct {
Login string `json:"login"` //auth login
//system level security's info
- Roles []string `json:"roles"`
+ Roles []string `json:"roles"`
//private fields
UserID string `json:"userid"` //system level user ID
@@ -83,7 +83,7 @@ func (u *UserSessionInfo) MustGetUserID() string {
return u.UserID
}
- panic(errors.NewWithHTTPCode(400, "invalid user"))
+ panic(errors.NewWithHTTPCode(401, "invalid user"))
}
func (u *UserSessionInfo) IsValid() bool {
@@ -92,7 +92,7 @@ func (u *UserSessionInfo) IsValid() bool {
if global.Env().IsDebug {
log.Error(util.MustToJSON(u), u.UserID)
}
- panic(errors.NewWithHTTPCode(400, "invalid user"))
+ return false
}
return v
}
diff --git a/core/security/user_session_test.go b/core/security/user_session_test.go
new file mode 100644
index 000000000..e857865bb
--- /dev/null
+++ b/core/security/user_session_test.go
@@ -0,0 +1,77 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package security
+
+import (
+ "encoding/json"
+ "testing"
+
+ "github.com/golang-jwt/jwt/v4"
+)
+
+func TestUserClaimsMarshalUsesFrameworkFields(t *testing.T) {
+ claims := UserClaims{
+ RegisteredClaims: &jwt.RegisteredClaims{},
+ UserSessionInfo: &UserSessionInfo{
+ Provider: "native",
+ Login: "admin@example.org",
+ Roles: []string{RoleAdmin},
+ UserID: "user-1",
+ },
+ }
+
+ payload, err := json.Marshal(claims)
+ if err != nil {
+ t.Fatalf("marshal claims: %v", err)
+ }
+
+ var data map[string]any
+ if err := json.Unmarshal(payload, &data); err != nil {
+ t.Fatalf("unmarshal claims json: %v", err)
+ }
+
+ if data["login"] != "admin@example.org" {
+ t.Fatalf("expected login field, got %#v", data["login"])
+ }
+ if data["userid"] != "user-1" {
+ t.Fatalf("expected userid field, got %#v", data["userid"])
+ }
+ if _, exists := data["username"]; exists {
+ t.Fatalf("did not expect legacy username alias in claims: %s", payload)
+ }
+ if _, exists := data["user_id"]; exists {
+ t.Fatalf("did not expect legacy user_id alias in claims: %s", payload)
+ }
+}
+
+func TestUserSessionInfoIsValidReturnsFalseForIncompleteUser(t *testing.T) {
+ user := &UserSessionInfo{
+ Provider: "native",
+ Login: "admin@example.org",
+ }
+
+ if user.IsValid() {
+ t.Fatal("expected incomplete user session to be invalid")
+ }
+}
diff --git a/core/security/validate.go b/core/security/validate.go
index d3574f427..d6f68a3de 100644
--- a/core/security/validate.go
+++ b/core/security/validate.go
@@ -16,11 +16,8 @@ import (
"infini.sh/framework/core/errors"
)
-func byAuthorizationHeader(w http.ResponseWriter, r *http.Request) (claims *UserClaims, err error) {
- var (
- authorization = r.Header.Get("Authorization")
- ok bool
- )
+func parseUserClaimsFromAuthorizationHeader(authorization string) (claims *UserClaims, err error) {
+ var ok bool
if authorization == "" {
return nil, errors.Error("Authorization not found")
@@ -64,6 +61,20 @@ func byAuthorizationHeader(w http.ResponseWriter, r *http.Request) (claims *User
return claims, nil
}
+// ValidateAuthorizationHeader validates a bearer token header and returns the
+// decoded framework session information for callers that only have header access.
+func ValidateAuthorizationHeader(authorization string) (*UserSessionInfo, error) {
+ claims, err := parseUserClaimsFromAuthorizationHeader(authorization)
+ if err != nil {
+ return nil, err
+ }
+ return claims.UserSessionInfo, nil
+}
+
+func byAuthorizationHeader(w http.ResponseWriter, r *http.Request) (claims *UserClaims, err error) {
+ return parseUserClaimsFromAuthorizationHeader(r.Header.Get("Authorization"))
+}
+
func ValidateLogin(w http.ResponseWriter, r *http.Request) (session *UserSessionInfo, err error) {
var claims *UserClaims
diff --git a/core/security/validate_test.go b/core/security/validate_test.go
new file mode 100644
index 000000000..6ad896983
--- /dev/null
+++ b/core/security/validate_test.go
@@ -0,0 +1,95 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package security
+
+import (
+ "testing"
+ "time"
+
+ "github.com/golang-jwt/jwt/v4"
+)
+
+// Header-only callers in downstream apps need the same validation path as the
+// framework HTTP auth middleware while shared auth code is being adopted.
+func TestValidateAuthorizationHeader(t *testing.T) {
+ oldSecret := secretKey
+ secretKey = "test-framework-secret"
+ defer func() {
+ secretKey = oldSecret
+ }()
+
+ token := jwt.NewWithClaims(jwt.SigningMethodHS256, UserClaims{
+ RegisteredClaims: &jwt.RegisteredClaims{
+ ExpiresAt: jwt.NewNumericDate(time.Now().Add(time.Hour)),
+ },
+ UserSessionInfo: &UserSessionInfo{
+ Provider: "native",
+ Login: "admin@example.org",
+ Roles: []string{RoleAdmin},
+ UserID: "user-1",
+ },
+ })
+ tokenString, err := token.SignedString([]byte(secretKey))
+ if err != nil {
+ t.Fatalf("sign token: %v", err)
+ }
+
+ sessionUser, err := ValidateAuthorizationHeader("Bearer " + tokenString)
+ if err != nil {
+ t.Fatalf("validate authorization header: %v", err)
+ }
+ if sessionUser.Login != "admin@example.org" {
+ t.Fatalf("expected login to round-trip, got %q", sessionUser.Login)
+ }
+ if sessionUser.UserID != "user-1" {
+ t.Fatalf("expected user id to round-trip, got %q", sessionUser.UserID)
+ }
+}
+
+func TestValidateAuthorizationHeaderRejectsIncompleteUserClaims(t *testing.T) {
+ oldSecret := secretKey
+ secretKey = "test-framework-secret"
+ defer func() {
+ secretKey = oldSecret
+ }()
+
+ token := jwt.NewWithClaims(jwt.SigningMethodHS256, UserClaims{
+ RegisteredClaims: &jwt.RegisteredClaims{
+ ExpiresAt: jwt.NewNumericDate(time.Now().Add(time.Hour)),
+ },
+ UserSessionInfo: &UserSessionInfo{
+ Provider: "native",
+ Login: "admin@example.org",
+ },
+ })
+ tokenString, err := token.SignedString([]byte(secretKey))
+ if err != nil {
+ t.Fatalf("sign token: %v", err)
+ }
+
+ sessionUser, err := ValidateAuthorizationHeader("Bearer " + tokenString)
+ if err == nil {
+ t.Fatalf("expected invalid claims to be rejected, got user %+v", sessionUser)
+ }
+}
diff --git a/core/task/chrono/task.go b/core/task/chrono/task.go
index d0b6aa8f8..b626b241d 100755
--- a/core/task/chrono/task.go
+++ b/core/task/chrono/task.go
@@ -35,9 +35,10 @@ import (
type Task func(ctx context.Context)
type SchedulerTask struct {
- task Task
- startTime time.Time
- location *time.Location
+ task Task
+ startTime time.Time
+ initialDelay time.Duration
+ location *time.Location
}
func CreateSchedulerTask(task Task, options ...Option) (*SchedulerTask, error) {
@@ -46,9 +47,10 @@ func CreateSchedulerTask(task Task, options ...Option) (*SchedulerTask, error) {
}
runnableTask := &SchedulerTask{
- task: task,
- startTime: time.Time{},
- location: time.Local,
+ task: task,
+ startTime: time.Time{},
+ initialDelay: 0,
+ location: time.Local,
}
for _, option := range options {
@@ -63,6 +65,10 @@ func CreateSchedulerTask(task Task, options ...Option) (*SchedulerTask, error) {
}
func (task *SchedulerTask) GetInitialDelay() time.Duration {
+ if task.initialDelay > 0 {
+ return task.initialDelay
+ }
+
if task.startTime.IsZero() {
return 0
}
@@ -87,6 +93,16 @@ func WithStartTime(year int, month time.Month, day, hour, min, sec int) Option {
}
}
+func WithInitialDelay(delay time.Duration) Option {
+ return func(task *SchedulerTask) error {
+ if delay < 0 {
+ delay = 0
+ }
+ task.initialDelay = delay
+ return nil
+ }
+}
+
func WithLocation(location string) Option {
return func(task *SchedulerTask) error {
loadedLocation, err := time.LoadLocation(location)
diff --git a/core/task/chrono/task_test.go b/core/task/chrono/task_test.go
index cba73252c..44f600b96 100755
--- a/core/task/chrono/task_test.go
+++ b/core/task/chrono/task_test.go
@@ -52,6 +52,16 @@ func TestNewSchedulerTask_WithInvalidLocation(t *testing.T) {
assert.Error(t, err)
}
+func TestNewSchedulerTask_WithInitialDelay(t *testing.T) {
+ task, err := CreateSchedulerTask(func(ctx context.Context) {
+ }, WithInitialDelay(200*time.Millisecond))
+ assert.Nil(t, err)
+
+ delay := task.GetInitialDelay()
+ assert.Greater(t, delay, 0*time.Millisecond)
+ assert.LessOrEqual(t, delay, 200*time.Millisecond)
+}
+
func TestNewScheduledRunnableTask(t *testing.T) {
task, _ := CreateScheduledRunnableTask(0, func(ctx context.Context) {
diff --git a/core/task/task.go b/core/task/task.go
index 8ae5d6d0e..a3c207e27 100644
--- a/core/task/task.go
+++ b/core/task/task.go
@@ -28,9 +28,11 @@ import (
log "github.com/cihub/seelog"
"infini.sh/framework/core/errors"
"infini.sh/framework/core/global"
+ "infini.sh/framework/core/orm"
"infini.sh/framework/core/task/chrono"
"infini.sh/framework/core/util"
"runtime"
+ "strings"
"sync"
"sync/atomic"
"time"
@@ -38,6 +40,22 @@ import (
var Tasks = sync.Map{}
+func shouldSilenceStartupTaskError(msg string) bool {
+ return !orm.HasHandler() && strings.Contains(msg, "ORM handler is not registered")
+}
+
+func logTaskRuntimeIssue(msg string, raw interface{}) {
+ if shouldSilenceStartupTaskError(msg) {
+ log.Debug(msg)
+ return
+ }
+ if raw != nil {
+ log.Error(raw, msg)
+ return
+ }
+ log.Error(msg)
+}
+
type State string
const (
@@ -103,7 +121,7 @@ func RegisterTransientTask(group, tag string, f func(ctx context.Context) error,
case string:
v = r.(string)
}
- log.Error(r, v)
+ logTaskRuntimeIssue(v, r)
}
}
task.State = Finished
@@ -118,7 +136,7 @@ func RegisterTransientTask(group, tag string, f func(ctx context.Context) error,
task.State = Running
err := inner(innerCtx)
if err != nil {
- log.Error(err)
+ logTaskRuntimeIssue(err.Error(), err)
}
t = time.Now()
task.EndTime = &t
@@ -128,15 +146,16 @@ func RegisterTransientTask(group, tag string, f func(ctx context.Context) error,
}
type ScheduleTask struct {
- ID string `config:"id" json:"id,omitempty"`
- Group string `config:"group" json:"group,omitempty"`
- Description string `config:"description" json:"description,omitempty"`
- Type string `config:"type" json:"type,omitempty"`
- Interval string `config:"interval" json:"interval,omitempty"`
- Crontab string `config:"crontab" json:"crontab,omitempty"`
- CreateTime time.Time `config:"create_time" json:"create_time,omitempty"`
- StartTime *time.Time `config:"start_time" json:"start_time,omitempty"`
- EndTime *time.Time `config:"end_time" json:"end_time,omitempty"`
+ ID string `config:"id" json:"id,omitempty"`
+ Group string `config:"group" json:"group,omitempty"`
+ Description string `config:"description" json:"description,omitempty"`
+ Type string `config:"type" json:"type,omitempty"`
+ Interval string `config:"interval" json:"interval,omitempty"`
+ InitialDelay string `config:"initial_delay" json:"initial_delay,omitempty"`
+ Crontab string `config:"crontab" json:"crontab,omitempty"`
+ CreateTime time.Time `config:"create_time" json:"create_time,omitempty"`
+ StartTime *time.Time `config:"start_time" json:"start_time,omitempty"`
+ EndTime *time.Time `config:"end_time" json:"end_time,omitempty"`
// Ensures the task runs as a singleton, preventing duplicate executions when previous attempt is not finished.
Singleton bool `config:"singleton" json:"singleton,omitempty"`
@@ -194,7 +213,7 @@ func RegisterScheduleTask(task ScheduleTask) (taskID string) {
case string:
v = r.(string)
}
- log.Error(v)
+ logTaskRuntimeIssue(v, nil)
}
}
task.isTaskRunning.Store(false)
@@ -232,6 +251,23 @@ var taskScheduler = chrono.NewDefaultTaskScheduler()
var defaultInterval = time.Duration(10) * time.Second
var started bool
+func getScheduleOptions(task *ScheduleTask) []chrono.Option {
+ if task == nil || task.Type != Interval || task.InitialDelay == "" {
+ return nil
+ }
+
+ initialDelay, err := time.ParseDuration(task.InitialDelay)
+ if err != nil {
+ log.Warnf("invalid initial delay for task [%s]: %s", task.ID, task.InitialDelay)
+ return nil
+ }
+ if initialDelay <= 0 {
+ return nil
+ }
+
+ return []chrono.Option{chrono.WithInitialDelay(initialDelay)}
+}
+
func RunTasks() {
started = true
Tasks.Range(func(key, value any) bool {
@@ -254,7 +290,7 @@ func runTask(task *ScheduleTask) {
switch task.Type {
case Interval:
- task1, err := taskScheduler.ScheduleAtFixedRate(task.Task, util.GetDurationOrDefault(task.Interval, defaultInterval))
+ task1, err := taskScheduler.ScheduleAtFixedRate(task.Task, util.GetDurationOrDefault(task.Interval, defaultInterval), getScheduleOptions(task)...)
if err != nil {
log.Error("failed to scheduled interval task:", task.Type, ",", task.Interval, ",", task.Description)
}
diff --git a/core/task/task_test.go b/core/task/task_test.go
new file mode 100644
index 000000000..678d9d88a
--- /dev/null
+++ b/core/task/task_test.go
@@ -0,0 +1,27 @@
+package task
+
+import "testing"
+
+func TestGetScheduleOptionsWithInitialDelay(t *testing.T) {
+ task := &ScheduleTask{
+ Type: Interval,
+ InitialDelay: "250ms",
+ }
+
+ options := getScheduleOptions(task)
+ if len(options) != 1 {
+ t.Fatalf("expected one schedule option, got %d", len(options))
+ }
+}
+
+func TestGetScheduleOptionsSkipsInvalidDelay(t *testing.T) {
+ task := &ScheduleTask{
+ Type: Interval,
+ InitialDelay: "invalid",
+ }
+
+ options := getScheduleOptions(task)
+ if len(options) != 0 {
+ t.Fatalf("expected no schedule options for invalid delay, got %d", len(options))
+ }
+}
diff --git a/core/util/fsutils.go b/core/util/fsutils.go
index 3e20aa483..6e9733e55 100755
--- a/core/util/fsutils.go
+++ b/core/util/fsutils.go
@@ -302,11 +302,11 @@ func FileExtension(file string) string {
return strings.ToLower(strings.TrimSpace(ext))
}
-// Smart get file abs path.
+// GetFileAbsPath resolves filePath to an absolute path when the file exists.
//
-// If all attempts fail, and `ignoreMissing` is set to `true`, this function
-// returns `filePath` as-is. Otherwise, it panics.
-func TryGetFileAbsPath(filePath string, ignoreMissing bool) string {
+// If all attempts fail and ignoreMissing is true, it returns filePath as-is.
+// Otherwise it returns an error describing the attempted paths.
+func GetFileAbsPath(filePath string, ignoreMissing bool) (string, error) {
// The paths that we tried
attempts := []string{}
@@ -317,7 +317,7 @@ func TryGetFileAbsPath(filePath string, ignoreMissing bool) string {
*/
if FileExists(filePath) {
- return filePath
+ return filePath, nil
} else {
attempts = append(attempts, filePath)
}
@@ -327,7 +327,7 @@ func TryGetFileAbsPath(filePath string, ignoreMissing bool) string {
*/
absPathRelativeToWd, _ := filepath.Abs(filePath)
if FileExists(absPathRelativeToWd) {
- return absPathRelativeToWd
+ return absPathRelativeToWd, nil
} else {
attempts = append(attempts, absPathRelativeToWd)
}
@@ -341,7 +341,7 @@ func TryGetFileAbsPath(filePath string, ignoreMissing bool) string {
absPathRelativeToExeDir := path.Join(exeDir, filePath)
if FileExists(absPathRelativeToExeDir) {
- return absPathRelativeToExeDir
+ return absPathRelativeToExeDir, nil
} else {
attempts = append(attempts, absPathRelativeToExeDir)
}
@@ -349,15 +349,27 @@ func TryGetFileAbsPath(filePath string, ignoreMissing bool) string {
}
/*
- * All attempts failed. Panic if `ignoreMissing` is not set. Otherwise,
- * return `filePath` as-is.
+ * All attempts failed. Return an error if `ignoreMissing` is not set.
+ * Otherwise, return `filePath` as-is.
*/
if !ignoreMissing {
errorMsg := fmt.Sprintf("failed to absolutize path '%s', tried %v, but they do not exist", filePath, attempts)
- panic(errors.New(errorMsg))
+ return "", errors.New(errorMsg)
} else {
- return filePath
+ return filePath, nil
+ }
+}
+
+// Smart get file abs path.
+//
+// If all attempts fail, and `ignoreMissing` is set to `true`, this function
+// returns `filePath` as-is. Otherwise, it panics.
+func TryGetFileAbsPath(filePath string, ignoreMissing bool) string {
+ absPath, err := GetFileAbsPath(filePath, ignoreMissing)
+ if err != nil {
+ panic(err)
}
+ return absPath
}
func ListAllFiles(path string) ([]string, error) {
diff --git a/core/util/fsutils_test.go b/core/util/fsutils_test.go
index 5613590d8..1392e80ad 100755
--- a/core/util/fsutils_test.go
+++ b/core/util/fsutils_test.go
@@ -42,6 +42,7 @@ package util
import (
"fmt"
"github.com/stretchr/testify/assert"
+ "os"
"path"
"path/filepath"
"testing"
@@ -145,3 +146,40 @@ func TestNormalizeFolderPath(t *testing.T) {
})
}
}
+
+func TestGetFileAbsPathReturnsAbsolutePathForExistingFile(t *testing.T) {
+ tempDir := t.TempDir()
+ configPath := filepath.Join(tempDir, "console.yml")
+ err := os.WriteFile(configPath, []byte("name: console\n"), 0644)
+ assert.NoError(t, err)
+
+ resolvedPath, err := GetFileAbsPath(configPath, false)
+ assert.NoError(t, err)
+ assert.Equal(t, configPath, resolvedPath)
+}
+
+func TestGetFileAbsPathReturnsErrorForMissingFile(t *testing.T) {
+ missingPath := filepath.Join(t.TempDir(), "missing-console.yml")
+
+ resolvedPath, err := GetFileAbsPath(missingPath, false)
+ assert.Error(t, err)
+ assert.Empty(t, resolvedPath)
+ assert.Contains(t, err.Error(), "failed to absolutize path")
+ assert.Contains(t, err.Error(), missingPath)
+}
+
+func TestGetFileAbsPathReturnsOriginalPathWhenMissingIsIgnored(t *testing.T) {
+ missingPath := filepath.Join(t.TempDir(), "missing-console.yml")
+
+ resolvedPath, err := GetFileAbsPath(missingPath, true)
+ assert.NoError(t, err)
+ assert.Equal(t, missingPath, resolvedPath)
+}
+
+func TestTryGetFileAbsPathPanicsForMissingFile(t *testing.T) {
+ missingPath := filepath.Join(t.TempDir(), "missing-console.yml")
+
+ assert.Panics(t, func() {
+ TryGetFileAbsPath(missingPath, false)
+ })
+}
diff --git a/core/vfs/static.go b/core/vfs/static.go
index 170b02d66..0b8e75cb7 100755
--- a/core/vfs/static.go
+++ b/core/vfs/static.go
@@ -103,7 +103,7 @@ func (fs StaticFS) Open(name string) (http.File, error) {
}
}
- log.Debug("local file not found,", localFile)
+ log.Trace("local file not found,", localFile)
}
if fs.SkipVFS {
diff --git a/docs/content.en/docs/release-notes/_index.md b/docs/content.en/docs/release-notes/_index.md
index 772ece46e..4ed7fcf7f 100644
--- a/docs/content.en/docs/release-notes/_index.md
+++ b/docs/content.en/docs/release-notes/_index.md
@@ -27,6 +27,7 @@ Information about release notes of INFINI Framework is provided here.
- feat(client): support token-based authorization #288
- feat: add pluggable sink to host metrics collectors #288
- feat: add access_token to security #359
+- feat(security): add native account login challenge, replay protection, and secure transport helpers
- feat: smtp processor support parse dynamic content attachments from message #374
- feat: add static rule based authorization #375
- feat: allow to specify OS user when installing the service #380
diff --git a/go.sum b/go.sum
index 4b3c93783..9d0d51a96 100644
--- a/go.sum
+++ b/go.sum
@@ -1,6 +1,5 @@
cloud.google.com/go v0.16.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
code.pfad.fr/check v1.1.0 h1:GWvjdzhSEgHvEHe2uJujDcpmZoySKuHQNrZMfzfO0bE=
-code.pfad.fr/check v1.1.0/go.mod h1:NiUH13DtYsb7xp5wll0U4SXx7KhXQVCtRgdC96IPfoM=
github.com/Azure/go-ntlmssp v0.0.0-20200615164410-66371956d46c/go.mod h1:chxPXzSsl7ZWRAuOIE23GDNzjWuZquvFlgA8xmpunjU=
github.com/Azure/go-ntlmssp v0.1.0 h1:DjFo6YtWzNqNvQdrwEyr/e4nhU3vRiwenz5QX7sFz+A=
github.com/Azure/go-ntlmssp v0.1.0/go.mod h1:NYqdhxd/8aAct/s4qSYZEerdPuH1liG2/X9DiVTbhpk=
@@ -13,7 +12,6 @@ github.com/PuerkitoBio/urlesc v0.0.0-20160726150825-5bd2802263f2/go.mod h1:uGdko
github.com/RoaringBitmap/roaring v1.9.4 h1:yhEIoH4YezLYT04s1nHehNO64EKFTop/wBhxv2QzDdQ=
github.com/RoaringBitmap/roaring v1.9.4/go.mod h1:6AXUsoIEzDTFFQCe1RbGA6uFONMhvejWj5rqITANK90=
github.com/alexbrainman/sspi v0.0.0-20250919150558-7d374ff0d59e h1:4dAU9FXIyQktpoUAgOJK3OTFc/xug0PCXYCqU0FgDKI=
-github.com/alexbrainman/sspi v0.0.0-20250919150558-7d374ff0d59e/go.mod h1:cEWa1LVoE5KvSD9ONXsZrj0z6KqySlCCNKHlLzbqAt4=
github.com/andybalholm/brotli v1.1.1 h1:PR2pgnyFznKEugtsUo0xLdDop5SKXd5Qf5ysW+7XdTA=
github.com/andybalholm/brotli v1.1.1/go.mod h1:05ib4cKhjx3OQYUY22hTVd34Bc8upXjOLL2rKwwZBoA=
github.com/arl/statsviz v0.6.0 h1:jbW1QJkEYQkufd//4NDYRSNBpwJNrdzPahF7ZmoGdyE=
@@ -42,7 +40,6 @@ github.com/dgraph-io/ristretto v0.2.0/go.mod h1:8uBHCU/PBV4Ag0CJrP47b9Ofby5dqWNh
github.com/dgraph-io/ristretto/v2 v2.2.0 h1:bkY3XzJcXoMuELV8F+vS8kzNgicwQFAaGINAEJdWGOM=
github.com/dgraph-io/ristretto/v2 v2.2.0/go.mod h1:RZrm63UmcBAaYWC1DotLYBmTvgkrs0+XhBd7Npn7/zI=
github.com/dgryski/go-farm v0.0.0-20240924180020-3414d57e47da h1:aIftn67I1fkbMa512G+w+Pxci9hJPB8oMnkcP3iZF38=
-github.com/dgryski/go-farm v0.0.0-20240924180020-3414d57e47da/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw=
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
github.com/docker/spdystream v0.0.0-20160310174837-449fdfce4d96/go.mod h1:Qh8CwZgvJUkLughtfhJv5dyTYa91l1fOUCrgjqmcifM=
@@ -58,7 +55,6 @@ github.com/evanphx/json-patch v0.0.0-20200808040245-162e5629780b/go.mod h1:NAJj0
github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM=
github.com/fatih/color v1.18.0/go.mod h1:4FelSpRwEGDpQ12mAdzqdOukCy4u8WUtOY6lkT/6HfU=
github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
-github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0=
github.com/fsnotify/fsnotify v1.4.3-0.20170329110642-4da3e2cfbabc/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
github.com/fsnotify/fsnotify v1.9.0 h1:2Ml+OJNzbYCTzsxtv8vKSFD9PbJjmhYF14k/jKC7S9k=
@@ -72,8 +68,6 @@ github.com/go-asn1-ber/asn1-ber v1.5.8-0.20250403174932-29230038a667 h1:BP4M0CvQ
github.com/go-asn1-ber/asn1-ber v1.5.8-0.20250403174932-29230038a667/go.mod h1:hEBeB/ic+5LoWskz+yKT7vGhhPYkProFKoKdwZRWMe0=
github.com/go-ini/ini v1.67.0 h1:z6ZrTEZqSWOTyH2FlglNbNgARyHG8oLW9gMELqKr06A=
github.com/go-ini/ini v1.67.0/go.mod h1:ByCAeIL28uOIIG0E3PJtZPDL8WnHpFKFOtgjp+3Ies8=
-github.com/go-jose/go-jose/v4 v4.1.3 h1:CVLmWDhDVRa6Mi/IgCgaopNosCaHz7zrMeF9MlZRkrs=
-github.com/go-jose/go-jose/v4 v4.1.3/go.mod h1:x4oUasVrzR7071A4TnHLGSPpNOm2a21K9Kf04k1rs08=
github.com/go-ldap/ldap/v3 v3.2.4/go.mod h1:iYS1MdmrmceOJ1QOTnRXrIs7i3kloqtmGQjRvjKpyMg=
github.com/go-ldap/ldap/v3 v3.4.13 h1:+x1nG9h+MZN7h/lUi5Q3UZ0fJ1GyDQYbPvbuH38baDQ=
github.com/go-ldap/ldap/v3 v3.4.13/go.mod h1:LxsGZV6vbaK0sIvYfsv47rfh4ca0JXokCoKjZxsszv0=
@@ -150,26 +144,13 @@ github.com/gorilla/sessions v1.4.0/go.mod h1:FLWm50oby91+hl7p/wRxDth9bWSuk0qVL2e
github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
github.com/gregjones/httpcache v0.0.0-20170920190843-316c5e0ff04e/go.mod h1:FecbI9+v66THATjSRHfNgh1IVFe/9kFxbXtjV0ctIMA=
-github.com/hashicorp/go-uuid v1.0.3 h1:2gKiV6YVmrJ1i2CKKa9obLvRieoRGviZFL26PcT/Co8=
-github.com/hashicorp/go-uuid v1.0.3/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/bN7x4byOro=
github.com/hashicorp/go-version v1.7.0 h1:5tqGy27NaOTB8yJKUZELlFAS/LTKJkrmONwQKeRZfjY=
github.com/hashicorp/go-version v1.7.0/go.mod h1:fltr4n8CU8Ke44wwGCBoEymUuxUHl09ZGVZPK5anwXA=
github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
github.com/hashicorp/hcl v0.0.0-20170914154624-68e816d1c783/go.mod h1:oZtUIOe8dh44I2q6ScRibXws4Ajl+d+nod3AaR9vL5w=
github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
github.com/inconshreveable/log15 v0.0.0-20170622235902-74a0988b5f80/go.mod h1:cOaXtrgN4ScfRrD9Bre7U1thNq5RtJ8ZoP4iXVGRj6o=
-github.com/jcmturner/aescts/v2 v2.0.0 h1:9YKLH6ey7H4eDBXW8khjYslgyqG2xZikXP0EQFKrle8=
-github.com/jcmturner/aescts/v2 v2.0.0/go.mod h1:AiaICIRyfYg35RUkr8yESTqvSy7csK90qZ5xfvvsoNs=
-github.com/jcmturner/dnsutils/v2 v2.0.0 h1:lltnkeZGL0wILNvrNiVCR6Ro5PGU/SeBvVO/8c/iPbo=
-github.com/jcmturner/dnsutils/v2 v2.0.0/go.mod h1:b0TnjGOvI/n42bZa+hmXL+kFJZsFT7G4t3HTlQ184QM=
-github.com/jcmturner/gofork v1.7.6 h1:QH0l3hzAU1tfT3rZCnW5zXl+orbkNMMRGJfdJjHVETg=
-github.com/jcmturner/gofork v1.7.6/go.mod h1:1622LH6i/EZqLloHfE7IeZ0uEJwMSUyQ/nDd82IeqRo=
-github.com/jcmturner/goidentity/v6 v6.0.1 h1:VKnZd2oEIMorCTsFBnJWbExfNN7yZr3EhJAxwOkZg6o=
-github.com/jcmturner/goidentity/v6 v6.0.1/go.mod h1:X1YW3bgtvwAXju7V3LCIMpY0Gbxyjn/mY9zx4tFonSg=
github.com/jcmturner/gokrb5/v8 v8.4.4 h1:x1Sv4HaTpepFkXbt2IkL29DXRf8sOfZXo8eRKh687T8=
-github.com/jcmturner/gokrb5/v8 v8.4.4/go.mod h1:1btQEpgT6k+unzCwX1KdWMEwPPkkgBtP+F6aCACiMrs=
-github.com/jcmturner/rpc/v2 v2.0.3 h1:7FXXj8Ti1IaVFpSAziCZWNzbNuZmnvw/i6CqLNdWfZY=
-github.com/jcmturner/rpc/v2 v2.0.3/go.mod h1:VUJYCIDm3PVOEHw8sgt091/20OJjskO/YJki3ELg/Hc=
github.com/jessevdk/go-flags v1.4.0/go.mod h1:4FA24M0QyGHXBuZZK/XkWh8h0e1EYbRYJSGM75WSRxI=
github.com/jmoiron/jsonq v0.0.0-20150511023944-e874b168d07e h1:ZZCvgaRDZg1gC9/1xrsgaJzQUCQgniKtw0xjWywWAOE=
github.com/jmoiron/jsonq v0.0.0-20150511023944-e874b168d07e/go.mod h1:+rHyWac2R9oAZwFe1wGY2HBzFJJy++RHBg1cU23NkD8=
@@ -193,16 +174,9 @@ github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzh
github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0=
github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
github.com/kr/pretty v0.2.0/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
-github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
-github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
-github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
-github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
-github.com/letsencrypt/challtestsrv v1.4.2 h1:0ON3ldMhZyWlfVNYYpFuWRTmZNnyfiL9Hh5YzC3JVwU=
-github.com/letsencrypt/challtestsrv v1.4.2/go.mod h1:GhqMqcSoeGpYd5zX5TgwA6er/1MbWzx/o7yuuVya+Wk=
github.com/letsencrypt/pebble/v2 v2.10.0 h1:Wq6gYXlsY6ubqI3hhxsTzdyotvfdjFBxuwYqCLCnj/U=
-github.com/letsencrypt/pebble/v2 v2.10.0/go.mod h1:Sk8cmUIPcIdv2nINo+9PB4L+ZBhzY+F9A1a/h/xmWiQ=
github.com/libdns/libdns v1.1.1 h1:wPrHrXILoSHKWJKGd0EiAVmiJbFShguILTg9leS/P/U=
github.com/libdns/libdns v1.1.1/go.mod h1:4Bj9+5CQiNMVGf87wjX4CY3HQJypUHRuLvlsfsZqLWQ=
github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ81pIr0yLvtUWk2if982qA3F3QD6H4=
@@ -219,7 +193,6 @@ github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovk
github.com/mattn/go-isatty v0.0.2/go.mod h1:M+lRXTBqGeGNdLjl/ufCoiOlB5xdOkqRJdNxMWT7Zi4=
github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
-github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc=
github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
github.com/mholt/acmez/v3 v3.1.6 h1:eGVQNObP0pBN4sxqrXeg7MYqTOWyoiYpQqITVWlrevk=
@@ -252,17 +225,13 @@ github.com/ncruces/julianday v1.0.0 h1:fH0OKwa7NWvniGQtxdJRxAgkBMolni2BjDHaWTxqt
github.com/ncruces/julianday v1.0.0/go.mod h1:Dusn2KvZrrovOMJuOt0TNXL6tB7U2E8kvza5fFc9G7g=
github.com/nsqio/nsq v1.3.0 h1:v7NtyO844ieTIOCQEqQ7IUSSi1ImhgrTTto1rgIYGEU=
github.com/nsqio/nsq v1.3.0/go.mod h1:RxNr6UC0kSkNF44LnJrlN3U3CQnQGTXk+QKfSZLzqvc=
-github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE=
-github.com/nxadm/tail v1.4.8/go.mod h1:+ncqLTQzXmGhMZNUePPaPqPvBxHAIsmXswZKocGu+AU=
github.com/onsi/ginkgo v0.0.0-20170829012221-11459a886d9c/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
github.com/onsi/ginkgo v1.11.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE=
-github.com/onsi/ginkgo v1.16.5/go.mod h1:+E8gABHa3K6zRBolWtd+ROzc/U5bkGt0FwiG042wbpU=
github.com/onsi/gomega v0.0.0-20170829124025-dcabb60a477c/go.mod h1:C1qb7wdrVGGVU+Z6iS04AVkA3Q65CEZX59MT0QO5uiA=
github.com/onsi/gomega v1.7.0/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY=
github.com/onsi/gomega v1.35.1 h1:Cwbd75ZBPxFSuZ6T+rN/WCb/gOc6YgFBXLlZLhC7Ds4=
-github.com/onsi/gomega v1.35.1/go.mod h1:PvZbdDc8J6XJEpDK4HCuRBm8a6Fzp9/DmhC9C7yFlog=
github.com/pelletier/go-toml v1.0.1-0.20170904195809-1d6b12b7cb29/go.mod h1:5z9KED0ma1S8pY6P1sdut58dfprrGBbd/94hg7ilaic=
github.com/pelletier/go-toml/v2 v2.2.3 h1:YmeHyLY8mFWbdkNWwpr+qIL2bEqT0o95WSdkNHvL12M=
github.com/pelletier/go-toml/v2 v2.2.3/go.mod h1:MfCQTFTvCcUyyvvwm1+G6H/jORL20Xlb6rzQu9GuUkc=
@@ -281,7 +250,6 @@ github.com/r3labs/diff/v2 v2.15.1/go.mod h1:I8noH9Fc2fjSaMxqF3G2lhDdC0b+JXCfyx85
github.com/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY=
github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII=
-github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o=
github.com/rs/cors v1.11.1 h1:eU3gRzXLRK57F5rKMGMZURNdIG4EoAmX8k94r9wXWHA=
github.com/rs/cors v1.11.1/go.mod h1:XyqrcTp5zjWr1wsJ8PIRZssZ8b/WMcMf71DJnit4EMU=
github.com/rs/xid v1.6.0 h1:fV591PaemRlL6JfRxGDEPl69wICngIQ3shQtzfy2gxU=
@@ -357,11 +325,9 @@ github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9dec
github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0=
github.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0=
github.com/zeebo/assert v1.1.0 h1:hU1L1vLTHsnO8x8c9KAR5GmM5QscxHg5RNU5z5qbUWY=
-github.com/zeebo/assert v1.1.0/go.mod h1:Pq9JiuJQpG8JLJdtkwrJESF0Foym2/D9XMU5ciN/wJ0=
github.com/zeebo/blake3 v0.2.4 h1:KYQPkhpRtcqh0ssGYcKLG1JYvddkEA8QwCM/yBqhaZI=
github.com/zeebo/blake3 v0.2.4/go.mod h1:7eeQ6d2iXWRGF6npfaxl2CU+xy2Fjo2gxeyZGCRUjcE=
github.com/zeebo/pcg v1.0.1 h1:lyqfGeWiv4ahac6ttHs+I5hwtH/+1mrhlCtVNQM2kHo=
-github.com/zeebo/pcg v1.0.1/go.mod h1:09F0S9iiKrwn9rlI5yjLkmrug154/YRW6KnnXVDM/l4=
github.com/zeebo/sbloom v0.0.0-20151106181526-405c65bd9be0 h1:EAluI/s9FYrMnDGmyXB6eKkjSNyn7lmSdvX975YHZnY=
github.com/zeebo/sbloom v0.0.0-20151106181526-405c65bd9be0/go.mod h1:J0OA/x7vNUsWZ88/oJ0BPtebbGfjvSW1lA07GinZNLM=
go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA=
@@ -373,7 +339,6 @@ go.opentelemetry.io/otel/metric v1.35.0/go.mod h1:nKVFgxBZ2fReX6IlyW28MgZojkoAkJ
go.opentelemetry.io/otel/trace v1.35.0 h1:dPpEfJu1sDIqruz7BHFG3c7528f6ddfSWfFDVt/xgMs=
go.opentelemetry.io/otel/trace v1.35.0/go.mod h1:WUk7DtFp1Aw2MkvqGdwiXYDZZNvA/1J8o6xRXLrIkyc=
go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
-go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0=
go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
go.uber.org/zap v1.27.1 h1:08RqriUEv8+ArZRYSTXy1LeBScaMpVSTBhCeaZYfMYc=
@@ -459,7 +424,6 @@ gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8
gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
-gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
gopkg.in/cheggaaa/pb.v1 v1.0.28 h1:n1tBJnnK2r7g9OW2btFH91V92STTUevLXYFb8gy9EMk=
gopkg.in/cheggaaa/pb.v1 v1.0.28/go.mod h1:V/YB90LKu/1FcN3WVnfiiE5oMCibMjukxqG/qStrOgw=
gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys=
@@ -470,7 +434,6 @@ gopkg.in/inf.v0 v0.9.1 h1:73M5CoZyi3ZLMOyDlQh031Cx6N9NDJ2Vvfl76EDAgDc=
gopkg.in/inf.v0 v0.9.1/go.mod h1:cWUDdTG/fYaXco+Dcufb5Vnc6Gp2YChqWtbxRZE0mXw=
gopkg.in/square/go-jose.v2 v2.6.0 h1:NGk74WTnPKBNUhNzQX7PYcTLUjoq7mzKk2OKbvwk2iI=
gopkg.in/square/go-jose.v2 v2.6.0/go.mod h1:M9dMgbHiYLoDGQrXy7OpJDJWiKiU//h+vD76mk0e1AI=
-gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ=
gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw=
gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
diff --git a/modules/api/api.go b/modules/api/api.go
index 67c7c01b3..e88cc4133 100755
--- a/modules/api/api.go
+++ b/modules/api/api.go
@@ -187,7 +187,9 @@ func (module *APIModule) Setup() {
}
func (module *APIModule) Start() error {
- api.StartAPI()
+ if global.Env().SystemConfig.APIConfig.Enabled {
+ api.StartAPI()
+ }
return nil
}
diff --git a/modules/configs/client/client.go b/modules/configs/client/client.go
index f863434b2..1a6b9e0ec 100644
--- a/modules/configs/client/client.go
+++ b/modules/configs/client/client.go
@@ -34,7 +34,9 @@ import (
"net/url"
"os"
"path/filepath"
+ "strings"
"sync"
+ "sync/atomic"
"time"
log "github.com/cihub/seelog"
@@ -46,72 +48,276 @@ import (
"infini.sh/framework/core/model"
"infini.sh/framework/core/task"
"infini.sh/framework/core/util"
+ ucfg "infini.sh/framework/lib/go-ucfg"
"infini.sh/framework/modules/configs/common"
"infini.sh/framework/modules/configs/config"
)
const bucketName = "instance_registered"
const configRegisterEnvKey = "CONFIG_MANAGED_SUCCESS"
+const legacyManagedRegisterCompatMaxVersion = "1.30.4"
+const unauthorizedRegisterRetryInterval = 10 * time.Second
+
+var postRegisterHooks []func(server string, res *util.Result) error
+var unauthorizedRegisterRetryLock sync.Mutex
+var lastUnauthorizedRegisterRetryAt time.Time
+var clearManagedRegistrationStateFunc = clearManagedRegistrationState
+var loadManagedBootstrapAccessTokenFunc = func() (string, error) {
+ return common.LoadTokenFromKeystore(common.ManagerBootstrapTokenKeystoreKey)
+}
+var restoreManagedBootstrapAccessTokenFunc = func() (string, error) {
+ token, err := loadManagedBootstrapAccessTokenFunc()
+ if err != nil {
+ return "", err
+ }
+ token = strings.TrimSpace(token)
+ if token == "" {
+ token = strings.TrimSpace(global.Env().SystemConfig.Configs.ManagerConfig.AccessToken.Get())
+ }
+ if token == "" {
+ token, err = common.LoadTokenFromKeystore(common.ManagerTokenKeystoreKey)
+ if err != nil {
+ return "", err
+ }
+ token = strings.TrimSpace(token)
+ }
+ if token == "" {
+ return "", fmt.Errorf("managed bootstrap access token is missing")
+ }
+ global.Env().SystemConfig.Configs.ManagerConfig.AccessToken = ucfg.SecretString(token)
+ return token, nil
+}
+var reconnectToManagerFunc func() error
+var configSyncInProgress atomic.Bool
-func ConnectToManager() error {
+func init() {
+ reconnectToManagerFunc = ConnectToManager
+}
+
+// maskURLInError replaces http(s):// URLs in error messages to avoid leaking internal addresses in logs.
+func maskURLInError(err error) string {
+ if err == nil {
+ return ""
+ }
+ msg := err.Error()
+ for _, scheme := range []string{"https://", "http://"} {
+ for {
+ idx := strings.Index(msg, scheme)
+ if idx < 0 {
+ break
+ }
+ end := strings.IndexAny(msg[idx:], " \"'\n\t")
+ if end < 0 {
+ msg = msg[:idx] + "***"
+ break
+ }
+ msg = msg[:idx] + "***" + msg[idx+end:]
+ }
+ }
+ return msg
+}
+
+func truncateManagerResponseBodyForLog(body []byte) string {
+ text := strings.TrimSpace(string(body))
+ if len(text) <= 256 {
+ return text
+ }
+ return text[:256] + "...(truncated)"
+}
+
+func tryStartManagedConfigSync() bool {
+ return configSyncInProgress.CompareAndSwap(false, true)
+}
+
+func finishManagedConfigSync() {
+ configSyncInProgress.Store(false)
+}
- if !global.Env().SystemConfig.Configs.Managed {
+func ConnectToManager() error {
+ cfg := global.Env().SystemConfig.Configs
+ if !cfg.Managed {
return nil
}
+ if cfg.Servers == nil || len(cfg.Servers) == 0 {
+ return errors.Errorf("no config manager was found")
+ }
// k8s env setting always_register_after_restart and pod after restart the ip will change so need register again
- if !global.Env().SystemConfig.Configs.AlwaysRegisterAfterRestart {
+ if !cfg.AlwaysRegisterAfterRestart {
if exists, err := kv.ExistsKey(bucketName, []byte(global.Env().SystemConfig.NodeConfig.ID)); exists && err == nil {
//already registered skip further process
- log.Info("already registered to config manager")
+ log.Infof("skip config manager registration for instance %v: local registration marker exists", global.Env().SystemConfig.NodeConfig.ID)
global.Register(configRegisterEnvKey, true)
return nil
}
}
- log.Info("register new instance to config manager")
-
- //register to config manager
- if global.Env().SystemConfig.Configs.Servers == nil || len(global.Env().SystemConfig.Configs.Servers) == 0 {
- return errors.Errorf("no config manager was found")
- }
-
info := model.GetInstanceInfo()
+ log.Infof("start config manager registration for instance %v against %d server(s)", info.ID, len(cfg.Servers))
+ registerReq := common.InstanceRegisterRequest{
+ Client: info,
+ }
+ registerAccessToken, err := buildManagedRegisterAccessToken(info)
+ if err != nil {
+ return err
+ }
+ if registerAccessToken != nil {
+ registerReq.AccessToken = registerAccessToken
+ }
req := util.Request{Method: util.Verb_POST}
req.ContentType = "application/json"
req.Path = common.REGISTER_API
- req.Body = util.MustToJSONBytes(info)
+ req.Body = util.MustToJSONBytes(registerReq)
server, res, err := submitRequestToManager(&req)
if err == nil && server != "" {
if res.StatusCode == 200 || util.ContainStr(string(res.Body), "exists") {
- log.Infof("success register to config manager: %v", string(server))
+ if err := execPostRegisterHooks(server, res); err != nil {
+ return err
+ }
+ log.Infof("config manager registration succeeded for instance %v via %v: status=%d", info.ID, server, res.StatusCode)
err := kv.AddValue(bucketName, []byte(global.Env().SystemConfig.NodeConfig.ID), []byte(util.GetLowPrecisionCurrentTime().String()))
if err != nil {
panic(err)
}
global.Register(configRegisterEnvKey, true)
+ } else {
+ if res.StatusCode == http.StatusUnauthorized {
+ if !claimUnauthorizedRegisterRetrySlot() {
+ return fmt.Errorf("unauthorized config manager registration")
+ }
+ return recoverManagedRegistrationWithBootstrap()
+ }
+ log.Warnf("config manager registration failed for instance %v via %v: status=%d, body=%s", info.ID, server, res.StatusCode, truncateManagerResponseBodyForLog(res.Body))
+ return fmt.Errorf("failed to register to config manager: status=%d, body=%s", res.StatusCode, strings.TrimSpace(string(res.Body)))
}
} else {
- log.Error("failed to register to config manager,", err, ",", server)
+ log.Errorf("config manager registration request failed for instance %v via %v: %v", info.ID, server, err)
}
return err
}
+func buildManagedRegisterAccessToken(info model.Instance) (*common.RegisterToken, error) {
+ if !common.SupportsManagedAccessToken(info.Application.Name) {
+ return nil, nil
+ }
+ if shouldSkipManagedRegisterAccessToken(info.Application.Version.VersionNumber) {
+ return nil, nil
+ }
+ accessToken, err := common.EnsureTokenInKeystore(common.AgentAccessTokenKeystoreKey)
+ if err != nil {
+ return nil, err
+ }
+ productName := strings.TrimSpace(info.Application.Name)
+ if productName == "" {
+ productName = "instance"
+ }
+ return &common.RegisterToken{
+ Name: fmt.Sprintf("%s access token", info.ID),
+ Description: fmt.Sprintf("Console to %s access token for instance %s", productName, info.ID),
+ Value: accessToken,
+ }, nil
+}
+
+func shouldSkipManagedRegisterAccessToken(version string) bool {
+ version = strings.TrimSpace(version)
+ if version == "" {
+ return false
+ }
+ parsed, err := util.ParseSemantic(version)
+ if err != nil {
+ parsed, err = util.ParseGeneric(version)
+ if err != nil {
+ return false
+ }
+ }
+ cmp, err := parsed.Compare(legacyManagedRegisterCompatMaxVersion)
+ if err != nil {
+ return false
+ }
+ return cmp <= 0
+}
+
+func AddPostRegisterHook(hook func(server string, res *util.Result) error) {
+ if hook != nil {
+ postRegisterHooks = append(postRegisterHooks, hook)
+ }
+}
+
+func clearManagedRegistrationState() error {
+ global.Register(configRegisterEnvKey, false)
+ instanceID := strings.TrimSpace(global.Env().SystemConfig.NodeConfig.ID)
+ if instanceID == "" {
+ return nil
+ }
+ return kv.DeleteKey(bucketName, []byte(instanceID))
+}
+
+func handleUnauthorizedConfigSyncResponse(res *util.Result) bool {
+ if res == nil || res.StatusCode != http.StatusUnauthorized {
+ return false
+ }
+
+ if !claimUnauthorizedRegisterRetrySlot() {
+ return true
+ }
+
+ log.Warn("config sync unauthorized, clearing local registration state and retrying registration")
+ if err := recoverManagedRegistrationWithBootstrap(); err != nil {
+ log.Warnf("failed to re-register to config manager after unauthorized config sync: %v", err)
+ return true
+ }
+ log.Info("re-registered to config manager after unauthorized config sync")
+ return true
+}
+
+func claimUnauthorizedRegisterRetrySlot() bool {
+ unauthorizedRegisterRetryLock.Lock()
+ defer unauthorizedRegisterRetryLock.Unlock()
+ if !lastUnauthorizedRegisterRetryAt.IsZero() && time.Since(lastUnauthorizedRegisterRetryAt) < unauthorizedRegisterRetryInterval {
+ return false
+ }
+ lastUnauthorizedRegisterRetryAt = time.Now()
+ return true
+}
+
+func recoverManagedRegistrationWithBootstrap() error {
+ if _, err := restoreManagedBootstrapAccessTokenFunc(); err != nil {
+ return err
+ }
+ if err := clearManagedRegistrationStateFunc(); err != nil {
+ return err
+ }
+ return reconnectToManagerFunc()
+}
+
+func execPostRegisterHooks(server string, res *util.Result) error {
+ for _, hook := range postRegisterHooks {
+ if err := hook(server, res); err != nil {
+ return err
+ }
+ }
+ return nil
+}
+
func submitRequestToManager(req *util.Request) (string, *util.Result, error) {
+ return DoManagerRequest(req)
+}
+
+func DoManagerRequest(req *util.Request) (string, *util.Result, error) {
var err error
var res *util.Result
cfg := global.Env().SystemConfig.Configs
- if cfg.ManagerConfig.BasicAuth.Username != "" {
- req.SetBasicAuth(cfg.ManagerConfig.BasicAuth.Username, cfg.ManagerConfig.BasicAuth.Password.Get())
+ if err = applyManagerRequestAuth(req); err != nil {
+ return "", nil, err
}
for _, server := range cfg.Servers {
req.Url, err = url.JoinPath(server, req.Path)
if err != nil {
continue
}
- res, err = util.ExecuteRequestWithCatchFlag(mTLSClient, req, true)
+ res, err = util.ExecuteRequestWithCatchFlag(getManagerHTTPClient(), req, true)
if err != nil {
continue
}
@@ -120,32 +326,72 @@ func submitRequestToManager(req *util.Request) (string, *util.Result, error) {
return "", nil, err
}
-var clientInitLock = sync.Once{}
+func applyManagerRequestAuth(req *util.Request) error {
+ cfg := global.Env().SystemConfig.Configs
+ if token := cfg.ManagerConfig.AccessToken.Get(); token != "" {
+ req.AddHeader(model.API_TOKEN, token)
+ return nil
+ }
+ token, err := common.LoadTokenFromKeystore(common.ManagerTokenKeystoreKey)
+ if err != nil {
+ return err
+ }
+ if token != "" {
+ req.AddHeader("Authorization", "Bearer "+token)
+ return nil
+ }
+ if cfg.ManagerConfig.BasicAuth.Username != "" {
+ req.SetBasicAuth(cfg.ManagerConfig.BasicAuth.Username, cfg.ManagerConfig.BasicAuth.Password.Get())
+ }
+ return nil
+}
+
+var managerHTTPClientInitLock = sync.Once{}
+var configSyncInitLock = sync.Once{}
var mTLSClient *http.Client
-func ListenConfigChanges() error {
+func initManagerHTTPClient() {
+ managerHTTPClientInitLock.Do(func() {
+ if !global.Env().SystemConfig.Configs.Managed {
+ return
+ }
+ cfg := global.Env().GetHTTPClientConfig("configs", "")
+ if cfg != nil {
+ hClient, err := api.NewHTTPClient(cfg)
+ if err != nil {
+ panic(err)
+ }
+ mTLSClient = hClient
+ }
+ })
+}
- clientInitLock.Do(func() {
+func getManagerHTTPClient() *http.Client {
+ initManagerHTTPClient()
+ return mTLSClient
+}
- if global.Env().SystemConfig.Configs.Managed {
- cfg := global.Env().GetHTTPClientConfig("configs", "")
- if cfg != nil {
- hClient, err := api.NewHTTPClient(cfg)
- if err != nil {
- panic(err)
- }
- mTLSClient = hClient
- }
+func ListenConfigChanges() error {
+ configSyncInitLock.Do(func() {
- //init config sync listening
- req := common.ConfigSyncRequest{}
- req.Client = model.GetInstanceInfo()
+ if global.Env().SystemConfig.Configs.Managed {
+ initManagerHTTPClient()
var syncFunc = func() {
+ if !tryStartManagedConfigSync() {
+ if global.Env().IsDebug {
+ log.Trace("skip overlapping config sync")
+ }
+ return
+ }
+ defer finishManagedConfigSync()
+
if global.Env().IsDebug {
log.Trace("fetch configs from manger")
}
+ req := common.ConfigSyncRequest{}
+ req.Client = model.GetInstanceInfo()
cfgs := config.GetConfigs(false, false)
req.Configs = cfgs
req.Hash = util.MD5digestString(util.MustToJSONBytes(cfgs))
@@ -154,19 +400,24 @@ func ListenConfigChanges() error {
request := util.Request{Method: util.Verb_POST}
request.ContentType = "application/json"
request.Path = common.SYNC_API
- request.Body = util.MustToJSONBytes(req)
+ requestBody := util.MustToJSONBytes(req)
+ request.Body = requestBody
if global.Env().IsDebug {
- log.Debug("config sync request: ", string(util.MustToJSONBytes(req)))
+ log.Debug("config sync request: ", string(requestBody))
}
- _, res, err := submitRequestToManager(&request)
+ _, res, err := DoManagerRequest(&request)
if err != nil {
- log.Error("failed to submit request to config manager,", err)
+ log.Error("failed to submit request to config manager,", maskURLInError(err))
return
}
if res != nil {
+ if handleUnauthorizedConfigSyncResponse(res) {
+ return
+ }
+
obj := common.ConfigSyncResponse{}
err := util.FromJSONBytes(res.Body, &obj)
if err != nil {
diff --git a/modules/configs/client/client_test.go b/modules/configs/client/client_test.go
new file mode 100644
index 000000000..965417018
--- /dev/null
+++ b/modules/configs/client/client_test.go
@@ -0,0 +1,301 @@
+package client
+
+import (
+ "net/http"
+ "net/http/httptest"
+ "os"
+ "path/filepath"
+ "strings"
+ "sync"
+ "sync/atomic"
+ "testing"
+ "time"
+
+ "infini.sh/framework/core/config"
+ "infini.sh/framework/core/env"
+ "infini.sh/framework/core/global"
+ "infini.sh/framework/core/model"
+ "infini.sh/framework/core/util"
+ ucfg "infini.sh/framework/lib/go-ucfg"
+)
+
+func TestApplyManagerRequestAuthUsesAccessTokenHeader(t *testing.T) {
+ oldConfigs := global.Env().SystemConfig.Configs
+ t.Cleanup(func() {
+ global.Env().SystemConfig.Configs = oldConfigs
+ })
+
+ global.Env().SystemConfig.Configs = config.ConfigsConfig{
+ ManagerConfig: struct {
+ LocalConfigsRepoPath string `config:"local_configs_repo_path"`
+ BasicAuth config.BasicAuth `config:"basic_auth"`
+ AccessToken ucfg.SecretString `config:"access_token"`
+ }{
+ AccessToken: ucfg.SecretString("manager-api-token"),
+ BasicAuth: config.BasicAuth{
+ Username: "manager",
+ Password: ucfg.SecretString("secret"),
+ },
+ },
+ }
+
+ req := &util.Request{}
+ if err := applyManagerRequestAuth(req); err != nil {
+ t.Fatalf("expected nil error, got %v", err)
+ }
+
+ headers := req.AllHeaders()
+ if headers[model.API_TOKEN] != "manager-api-token" {
+ t.Fatalf("expected %s header to be set, got %#v", model.API_TOKEN, headers)
+ }
+ if auth := headers["Authorization"]; auth != "" {
+ t.Fatalf("expected no Authorization header, got %q", auth)
+ }
+}
+
+func TestBuildManagedRegisterAccessToken(t *testing.T) {
+ t.Setenv("KEYSTORE_PATH", t.TempDir())
+
+ instance := model.Instance{}
+ instance.ID = "gateway-1"
+ instance.Application = env.Application{
+ Name: "gateway",
+ Version: env.Version{VersionNumber: "1.30.5"},
+ }
+
+ registerToken, err := buildManagedRegisterAccessToken(instance)
+ if err != nil {
+ t.Fatalf("expected nil error, got %v", err)
+ }
+ if registerToken == nil || registerToken.Value == "" {
+ t.Fatalf("expected managed register token, got %#v", registerToken)
+ }
+ if !strings.Contains(registerToken.Description, "gateway") {
+ t.Fatalf("unexpected description: %q", registerToken.Description)
+ }
+
+ other := model.Instance{}
+ other.ID = "other-1"
+ other.Application = env.Application{Name: "console"}
+ registerToken, err = buildManagedRegisterAccessToken(other)
+ if err != nil {
+ t.Fatalf("expected nil error, got %v", err)
+ }
+ if registerToken != nil {
+ t.Fatalf("expected no managed register token, got %#v", registerToken)
+ }
+
+ legacy := model.Instance{}
+ legacy.ID = "legacy-agent-1"
+ legacy.Application = env.Application{
+ Name: "agent",
+ Version: env.Version{VersionNumber: "1.30.4"},
+ }
+ registerToken, err = buildManagedRegisterAccessToken(legacy)
+ if err != nil {
+ t.Fatalf("expected nil error for legacy agent, got %v", err)
+ }
+ if registerToken != nil {
+ t.Fatalf("expected legacy agent to skip managed register token, got %#v", registerToken)
+ }
+}
+
+func TestListenConfigChangesStillSyncsAfterHTTPClientInit(t *testing.T) {
+ var syncRequests atomic.Int32
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if r.URL.Path != "/configs/_sync" {
+ t.Fatalf("unexpected request path: %s", r.URL.Path)
+ }
+ if r.Method != http.MethodPost {
+ t.Fatalf("unexpected request method: %s", r.Method)
+ }
+ syncRequests.Add(1)
+ w.Header().Set("Content-Type", "application/json")
+ _, _ = w.Write([]byte(`{"changed":false}`))
+ }))
+ defer server.Close()
+
+ tempDir := t.TempDir()
+ configDir := filepath.Join(tempDir, "configs")
+ if err := os.MkdirAll(configDir, 0o755); err != nil {
+ t.Fatalf("create config dir: %v", err)
+ }
+ mainConfigFile := filepath.Join(tempDir, "agent.yml")
+ if err := os.WriteFile(mainConfigFile, []byte("node:\n id: test-agent\n"), 0o644); err != nil {
+ t.Fatalf("write main config: %v", err)
+ }
+
+ oldEnv := global.Env()
+ oldHTTPClientInitLock := managerHTTPClientInitLock
+ oldConfigSyncInitLock := configSyncInitLock
+ oldClient := mTLSClient
+ t.Cleanup(func() {
+ global.RegisterEnv(oldEnv)
+ managerHTTPClientInitLock = oldHTTPClientInitLock
+ configSyncInitLock = oldConfigSyncInitLock
+ mTLSClient = oldClient
+ })
+
+ testEnv := env.EmptyEnv()
+ testEnv.SystemConfig.Configs.Managed = true
+ testEnv.SystemConfig.Configs.Servers = []string{server.URL}
+ testEnv.SystemConfig.Configs.Interval = "30s"
+ testEnv.SystemConfig.PathConfig.Config = configDir
+ testEnv.SystemConfig.NodeConfig.ID = "test-agent"
+ testEnv.SetConfigFile(mainConfigFile)
+ global.RegisterEnv(testEnv)
+
+ managerHTTPClientInitLock = sync.Once{}
+ configSyncInitLock = sync.Once{}
+ mTLSClient = nil
+
+ if getManagerHTTPClient() == nil {
+ t.Fatal("expected manager HTTP client to initialize")
+ }
+ if err := ListenConfigChanges(); err != nil {
+ t.Fatalf("expected nil error, got %v", err)
+ }
+ if syncRequests.Load() != 1 {
+ t.Fatalf("expected one immediate sync request, got %d", syncRequests.Load())
+ }
+}
+
+func TestHandleUnauthorizedConfigSyncResponseClearsStateAndReconnects(t *testing.T) {
+ oldClear := clearManagedRegistrationStateFunc
+ oldReconnect := reconnectToManagerFunc
+ oldLoadBootstrap := loadManagedBootstrapAccessTokenFunc
+ oldRestoreBootstrap := restoreManagedBootstrapAccessTokenFunc
+ oldRetryAt := lastUnauthorizedRegisterRetryAt
+ oldAccessToken := global.Env().SystemConfig.Configs.ManagerConfig.AccessToken
+ t.Cleanup(func() {
+ clearManagedRegistrationStateFunc = oldClear
+ reconnectToManagerFunc = oldReconnect
+ loadManagedBootstrapAccessTokenFunc = oldLoadBootstrap
+ restoreManagedBootstrapAccessTokenFunc = oldRestoreBootstrap
+ lastUnauthorizedRegisterRetryAt = oldRetryAt
+ global.Env().SystemConfig.Configs.ManagerConfig.AccessToken = oldAccessToken
+ })
+
+ var cleared atomic.Int32
+ var reconnected atomic.Int32
+ var restored atomic.Int32
+ clearManagedRegistrationStateFunc = func() error {
+ cleared.Add(1)
+ return nil
+ }
+ reconnectToManagerFunc = func() error {
+ reconnected.Add(1)
+ return nil
+ }
+ loadManagedBootstrapAccessTokenFunc = func() (string, error) {
+ restored.Add(1)
+ return "bootstrap-token", nil
+ }
+ restoreManagedBootstrapAccessTokenFunc = func() (string, error) {
+ token, err := loadManagedBootstrapAccessTokenFunc()
+ if err != nil {
+ return "", err
+ }
+ global.Env().SystemConfig.Configs.ManagerConfig.AccessToken = ucfg.SecretString(token)
+ return token, nil
+ }
+ lastUnauthorizedRegisterRetryAt = time.Time{}
+ global.Env().SystemConfig.Configs.ManagerConfig.AccessToken = ""
+
+ handled := handleUnauthorizedConfigSyncResponse(&util.Result{StatusCode: http.StatusUnauthorized})
+ if !handled {
+ t.Fatal("expected unauthorized config sync response to be handled")
+ }
+ if restored.Load() != 1 {
+ t.Fatalf("expected bootstrap token to be loaded once, got %d", restored.Load())
+ }
+ if cleared.Load() != 1 {
+ t.Fatalf("expected local registration state to be cleared once, got %d", cleared.Load())
+ }
+ if reconnected.Load() != 1 {
+ t.Fatalf("expected reconnect to run once, got %d", reconnected.Load())
+ }
+ if got := global.Env().SystemConfig.Configs.ManagerConfig.AccessToken.Get(); got != "bootstrap-token" {
+ t.Fatalf("expected bootstrap token to be restored, got %q", got)
+ }
+
+ handled = handleUnauthorizedConfigSyncResponse(&util.Result{StatusCode: http.StatusUnauthorized})
+ if !handled {
+ t.Fatal("expected throttled unauthorized config sync response to still be handled")
+ }
+ if cleared.Load() != 1 {
+ t.Fatalf("expected throttled retry not to clear state again, got %d", cleared.Load())
+ }
+ if reconnected.Load() != 1 {
+ t.Fatalf("expected throttled retry not to reconnect again, got %d", reconnected.Load())
+ }
+}
+
+func TestManagedConfigSyncGuardPreventsOverlap(t *testing.T) {
+ configSyncInProgress.Store(false)
+ t.Cleanup(func() {
+ configSyncInProgress.Store(false)
+ })
+
+ if !tryStartManagedConfigSync() {
+ t.Fatal("expected first config sync to start")
+ }
+ if tryStartManagedConfigSync() {
+ t.Fatal("expected overlapping config sync to be rejected")
+ }
+
+ finishManagedConfigSync()
+
+ if !tryStartManagedConfigSync() {
+ t.Fatal("expected config sync to start again after previous one finished")
+ }
+ finishManagedConfigSync()
+}
+
+func TestRestoreManagedBootstrapAccessTokenFallsBackToManagerAccessToken(t *testing.T) {
+ t.Setenv("KEYSTORE_PATH", t.TempDir())
+
+ oldLoadBootstrap := loadManagedBootstrapAccessTokenFunc
+ oldAccessToken := global.Env().SystemConfig.Configs.ManagerConfig.AccessToken
+ t.Cleanup(func() {
+ loadManagedBootstrapAccessTokenFunc = oldLoadBootstrap
+ global.Env().SystemConfig.Configs.ManagerConfig.AccessToken = oldAccessToken
+ })
+
+ loadManagedBootstrapAccessTokenFunc = func() (string, error) {
+ return "", nil
+ }
+ global.Env().SystemConfig.Configs.ManagerConfig.AccessToken = ucfg.SecretString("config-access-token")
+
+ token, err := restoreManagedBootstrapAccessTokenFunc()
+ if err != nil {
+ t.Fatalf("expected nil error, got %v", err)
+ }
+ if token != "config-access-token" {
+ t.Fatalf("expected config access token fallback, got %q", token)
+ }
+}
+
+func TestRestoreManagedBootstrapAccessTokenReturnsErrorWhenNoFallbackAvailable(t *testing.T) {
+ t.Setenv("KEYSTORE_PATH", t.TempDir())
+
+ oldLoadBootstrap := loadManagedBootstrapAccessTokenFunc
+ oldAccessToken := global.Env().SystemConfig.Configs.ManagerConfig.AccessToken
+ t.Cleanup(func() {
+ loadManagedBootstrapAccessTokenFunc = oldLoadBootstrap
+ global.Env().SystemConfig.Configs.ManagerConfig.AccessToken = oldAccessToken
+ })
+
+ loadManagedBootstrapAccessTokenFunc = func() (string, error) {
+ return "", nil
+ }
+ global.Env().SystemConfig.Configs.ManagerConfig.AccessToken = ""
+
+ _, err := restoreManagedBootstrapAccessTokenFunc()
+ if err == nil {
+ t.Fatal("expected missing bootstrap token error")
+ }
+ if !strings.Contains(err.Error(), "managed bootstrap access token is missing") {
+ t.Fatalf("unexpected error: %v", err)
+ }
+}
diff --git a/modules/configs/common/config.go b/modules/configs/common/config.go
index 03ce8dbe4..a13c8675e 100644
--- a/modules/configs/common/config.go
+++ b/modules/configs/common/config.go
@@ -38,11 +38,12 @@ type AgentConfig struct {
}
type SetupConfig struct {
- DownloadURL string `config:"download_url"`
- CACertFile string `config:"ca_cert"`
- CAKeyFile string `config:"ca_key"`
- ConsoleEndpoint string `config:"console_endpoint"`
- Port string `config:"port"`
+ DownloadURL string `config:"download_url"`
+ CACertFile string `config:"ca_cert"`
+ CAKeyFile string `config:"ca_key"`
+ ConsoleEndpoint string `config:"console_endpoint"`
+ ReverseChannelEndpoints []string `config:"reverse_channel_endpoints"`
+ Port string `config:"port"`
}
func GetAgentConfig() *AgentConfig {
diff --git a/modules/configs/common/domain.go b/modules/configs/common/domain.go
index 232c1fc52..23faabba1 100644
--- a/modules/configs/common/domain.go
+++ b/modules/configs/common/domain.go
@@ -27,11 +27,32 @@
package common
-import "infini.sh/framework/core/model"
+import (
+ "strings"
+
+ "infini.sh/framework/core/model"
+)
const REGISTER_API = "/instance/_register"
const SYNC_API = "/configs/_sync"
+const (
+ ManagerTokenKeystoreKey = "configs_manager_token"
+ ManagerBootstrapTokenKeystoreKey = "configs_manager_bootstrap_token"
+ AgentAccessTokenKeystoreKey = "agent_access_token"
+)
+
+type RegisterToken struct {
+ Name string `json:"name,omitempty"`
+ Description string `json:"description,omitempty"`
+ Value string `json:"value,omitempty"`
+}
+
+type InstanceRegisterRequest struct {
+ Client model.Instance `json:"client"`
+ AccessToken *RegisterToken `json:"access_token,omitempty"`
+}
+
type ConfigFile struct {
Name string `json:"name,omitempty"`
Location string `json:"location,omitempty"`
@@ -109,3 +130,12 @@ type InstanceSettings struct {
ConfigFiles []string `config:"configs"`
Secrets []string `config:"secrets"`
}
+
+func SupportsManagedAccessToken(applicationName string) bool {
+ switch strings.ToLower(strings.TrimSpace(applicationName)) {
+ case "agent", "gateway":
+ return true
+ default:
+ return false
+ }
+}
diff --git a/modules/configs/common/token.go b/modules/configs/common/token.go
new file mode 100644
index 000000000..7f7e634cc
--- /dev/null
+++ b/modules/configs/common/token.go
@@ -0,0 +1,66 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+/* Copyright © INFINI LTD. All rights reserved.
+ * Web: https://infinilabs.com
+ * Email: hello#infini.ltd */
+
+package common
+
+import (
+ "strings"
+
+ "infini.sh/framework/core/keystore"
+ "infini.sh/framework/core/util"
+ keystore2 "infini.sh/framework/lib/keystore"
+)
+
+func LoadTokenFromKeystore(key string) (string, error) {
+ value, err := keystore.GetValue(key)
+ if err == keystore2.ErrKeyDoesntExists {
+ return "", nil
+ }
+ if err != nil {
+ return "", err
+ }
+ return strings.TrimSpace(string(value)), nil
+}
+
+func SaveTokenToKeystore(key, value string) error {
+ return keystore.SetValue(key, util.UnsafeStringToBytes(strings.TrimSpace(value)))
+}
+
+func EnsureTokenInKeystore(key string) (string, error) {
+ value, err := LoadTokenFromKeystore(key)
+ if err != nil {
+ return "", err
+ }
+ if value != "" {
+ return value, nil
+ }
+ value = util.GenerateRandomString(48)
+ if err := SaveTokenToKeystore(key, value); err != nil {
+ return "", err
+ }
+ return value, nil
+}
diff --git a/modules/elastic/adapter/elasticsearch/v0.go b/modules/elastic/adapter/elasticsearch/v0.go
index 0a13b8bd8..bd7b0f7c5 100755
--- a/modules/elastic/adapter/elasticsearch/v0.go
+++ b/modules/elastic/adapter/elasticsearch/v0.go
@@ -491,6 +491,20 @@ func (c *ESAPIV0) Get(indexName, docType, id string) (*elastic.GetResponse, erro
return esResp, err
}
+ if resp.StatusCode >= 400 {
+ if esResp.Error != nil {
+ errType := esResp.Error.Type
+ errReason := esResp.Error.Message()
+ if errType != "" && errReason != "" {
+ return esResp, errors.Errorf("status:%d, type:%s, reason:%s", resp.StatusCode, errType, errReason)
+ }
+ if errReason != "" {
+ return esResp, errors.Errorf("status:%d, reason:%s", resp.StatusCode, errReason)
+ }
+ }
+ return esResp, errors.Errorf("status:%d", resp.StatusCode)
+ }
+
return esResp, nil
}
@@ -577,7 +591,7 @@ func (c *ESAPIV0) Search(indexName string, query *elastic.SearchRequest) (*elast
js := query.ToJSONString()
if global.Env().IsDebug {
- log.Info(js)
+ log.Trace(js)
}
return c.SearchWithRawQueryDSL(indexName, util.UnsafeStringToBytes(js))
@@ -607,6 +621,18 @@ func (c *ESAPIV0) QueryDSL(ctx context.Context, indexName string, queryArgs *[]u
}
resp, err := c.Request(ctx, util.Verb_POST, url, queryDSL)
+ if err == nil && resp != nil && shouldRetryWithoutTermsMissing(resp.StatusCode, resp.Body) {
+ if retryQueryDSL, changed := stripTermsMissingFromQueryDSL(queryDSL); changed {
+ if global.Env().IsDebug {
+ log.Tracef("retrying query without terms.missing due to UnmappedTerms response: %s", url)
+ }
+ retryResp, retryErr := c.Request(ctx, util.Verb_POST, url, retryQueryDSL)
+ if retryErr == nil && retryResp != nil {
+ resp = retryResp
+ queryDSL = retryQueryDSL
+ }
+ }
+ }
if resp != nil {
esResp.StatusCode = resp.StatusCode
esResp.RawResult = resp
@@ -633,6 +659,61 @@ func (c *ESAPIV0) QueryDSL(ctx context.Context, indexName string, queryArgs *[]u
return esResp, nil
}
+func shouldRetryWithoutTermsMissing(statusCode int, body []byte) bool {
+ if statusCode < 500 || len(body) == 0 {
+ return false
+ }
+ lowerBody := strings.ToLower(util.UnsafeBytesToString(body))
+ return strings.Contains(lowerBody, "unmappedterms") &&
+ strings.Contains(lowerBody, "unsupported")
+}
+
+func stripTermsMissingFromQueryDSL(queryDSL []byte) ([]byte, bool) {
+ if len(queryDSL) == 0 {
+ return nil, false
+ }
+ payload := map[string]interface{}{}
+ if err := json.Unmarshal(queryDSL, &payload); err != nil {
+ return nil, false
+ }
+ changed := stripTermsMissingRecursive(payload)
+ if !changed {
+ return nil, false
+ }
+ newDSL, err := json.Marshal(payload)
+ if err != nil {
+ return nil, false
+ }
+ return newDSL, true
+}
+
+func stripTermsMissingRecursive(value interface{}) bool {
+ changed := false
+ switch typed := value.(type) {
+ case map[string]interface{}:
+ if termsValue, ok := typed["terms"]; ok {
+ if termsMap, ok := termsValue.(map[string]interface{}); ok {
+ if _, exists := termsMap["missing"]; exists {
+ delete(termsMap, "missing")
+ changed = true
+ }
+ }
+ }
+ for _, nested := range typed {
+ if stripTermsMissingRecursive(nested) {
+ changed = true
+ }
+ }
+ case []interface{}:
+ for _, nested := range typed {
+ if stripTermsMissingRecursive(nested) {
+ changed = true
+ }
+ }
+ }
+ return changed
+}
+
func (c *ESAPIV0) SearchWithRawQueryDSL(indexName string, queryDSL []byte) (*elastic.SearchResponse, error) {
return c.QueryDSL(nil, indexName, nil, queryDSL)
}
@@ -1517,6 +1598,9 @@ func (c *ESAPIV0) GetAliases() (*map[string]elastic.AliasInfo, error) {
resp, err := c.Request(nil, util.Verb_GET, url, nil)
if err != nil || resp.StatusCode != 200 {
+ if err == nil {
+ return nil, errors.NewWithHTTPCode(resp.StatusCode, string(resp.Body))
+ }
return nil, err
}
@@ -1599,6 +1683,9 @@ func (c *ESAPIV0) GetAliasesAndIndices() (*elastic.AliasAndIndicesResponse, erro
resp, err := c.Request(nil, util.Verb_GET, url, nil)
if err != nil || resp.StatusCode != 200 {
+ if err == nil {
+ return nil, errors.NewWithHTTPCode(resp.StatusCode, string(resp.Body))
+ }
return nil, err
}
data := map[string]AliasesResponse{}
diff --git a/modules/elastic/adapter/elasticsearch/v0_querydsl_test.go b/modules/elastic/adapter/elasticsearch/v0_querydsl_test.go
new file mode 100644
index 000000000..b5cc1d430
--- /dev/null
+++ b/modules/elastic/adapter/elasticsearch/v0_querydsl_test.go
@@ -0,0 +1,69 @@
+package elasticsearch
+
+import (
+ "testing"
+
+ "github.com/segmentio/encoding/json"
+)
+
+func TestStripTermsMissingFromQueryDSL(t *testing.T) {
+ source := []byte(`{
+ "aggs": {
+ "a": {
+ "terms": {
+ "field": "metadata.labels.cluster_id",
+ "missing": "",
+ "size": 2
+ },
+ "aggs": {
+ "b": {
+ "date_range": {
+ "field": "timestamp",
+ "ranges": [{"from":"now-1d/d","to":"now/d"}]
+ },
+ "aggs": {
+ "c": {
+ "terms": {
+ "field": "payload.elasticsearch.cluster_health.status",
+ "missing": "",
+ "size": 2
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }`)
+
+ got, changed := stripTermsMissingFromQueryDSL(source)
+ if !changed {
+ t.Fatal("expected query DSL to be changed")
+ }
+ var parsed map[string]interface{}
+ if err := json.Unmarshal(got, &parsed); err != nil {
+ t.Fatalf("expected valid JSON, got %v", err)
+ }
+
+ aggA := parsed["aggs"].(map[string]interface{})["a"].(map[string]interface{})
+ termsA := aggA["terms"].(map[string]interface{})
+ if _, ok := termsA["missing"]; ok {
+ t.Fatalf("expected top-level terms.missing to be removed, got %#v", termsA)
+ }
+
+ aggB := aggA["aggs"].(map[string]interface{})["b"].(map[string]interface{})
+ aggC := aggB["aggs"].(map[string]interface{})["c"].(map[string]interface{})
+ termsC := aggC["terms"].(map[string]interface{})
+ if _, ok := termsC["missing"]; ok {
+ t.Fatalf("expected nested terms.missing to be removed, got %#v", termsC)
+ }
+}
+
+func TestShouldRetryWithoutTermsMissing(t *testing.T) {
+ if shouldRetryWithoutTermsMissing(400, []byte(`{"error":{"reason":"UnmappedTerms unsupported"}}`)) {
+ t.Fatal("should not retry on non-5xx status")
+ }
+ if !shouldRetryWithoutTermsMissing(500, []byte(`{"error":{"reason":"Aggregation [x] is of type [UnmappedTerms] which is currently unsupported."}}`)) {
+ t.Fatal("expected retry to be enabled for UnmappedTerms unsupported error")
+ }
+}
diff --git a/modules/elastic/adapter/elasticsearch/v7.go b/modules/elastic/adapter/elasticsearch/v7.go
index 9b0c51c32..9648fbf70 100755
--- a/modules/elastic/adapter/elasticsearch/v7.go
+++ b/modules/elastic/adapter/elasticsearch/v7.go
@@ -374,7 +374,7 @@ func (c *ESAPIV7) Create(indexName, docType string, id interface{}, data interfa
}
if global.Env().IsDebug {
- log.Debug("creating doc: ", url, ",", string(js))
+ log.Trace("creating doc: ", url, ",", string(js))
}
if err != nil {
diff --git a/modules/elastic/adapter/elasticsearch/v8.go b/modules/elastic/adapter/elasticsearch/v8.go
index c6f0dc0a4..427be5bea 100644
--- a/modules/elastic/adapter/elasticsearch/v8.go
+++ b/modules/elastic/adapter/elasticsearch/v8.go
@@ -280,7 +280,7 @@ func (c *ESAPIV8) Create(indexName, docType string, id interface{}, data interfa
}
if global.Env().IsDebug {
- log.Debug("creating doc: ", url, ",", string(js))
+ log.Trace("creating doc: ", url, ",", string(js))
}
if err != nil {
diff --git a/modules/elastic/adapter/ver.go b/modules/elastic/adapter/ver.go
index 1e3120cb5..a02119217 100755
--- a/modules/elastic/adapter/ver.go
+++ b/modules/elastic/adapter/ver.go
@@ -169,6 +169,14 @@ func RequestTimeout(ctx *elastic.APIContext, method, url string, body []byte, me
func GetClusterUUID(clusterID string) (string, error) {
meta := elastic.GetMetadata(clusterID)
+ if meta == nil {
+ if cfg := elastic.GetConfigNoPanic(clusterID); cfg != nil {
+ if cfg.ClusterUUID != "" {
+ return cfg.ClusterUUID, nil
+ }
+ meta = elastic.GetOrInitMetadata(cfg)
+ }
+ }
if meta == nil {
return "", fmt.Errorf("metadata can not be mepty")
}
diff --git a/modules/elastic/adapter/ver_test.go b/modules/elastic/adapter/ver_test.go
new file mode 100644
index 000000000..db5808e55
--- /dev/null
+++ b/modules/elastic/adapter/ver_test.go
@@ -0,0 +1,30 @@
+package adapter
+
+import (
+ "testing"
+
+ "infini.sh/framework/core/elastic"
+ "infini.sh/framework/core/orm"
+)
+
+func TestGetClusterUUIDFallsBackToConfigWhenMetadataMissing(t *testing.T) {
+ cfg := elastic.ElasticsearchConfig{
+ ORMObjectBase: orm.ORMObjectBase{ID: "test-cluster-uuid-fallback"},
+ Name: "test-cluster-uuid-fallback",
+ ClusterUUID: "cluster-uuid-fallback",
+ }
+
+ t.Cleanup(func() {
+ elastic.RemoveInstance(cfg.ID)
+ })
+
+ elastic.UpdateConfig(cfg)
+
+ clusterUUID, err := GetClusterUUID(cfg.ID)
+ if err != nil {
+ t.Fatalf("expected cluster uuid from config fallback, got error: %v", err)
+ }
+ if clusterUUID != cfg.ClusterUUID {
+ t.Fatalf("expected cluster uuid %q, got %q", cfg.ClusterUUID, clusterUUID)
+ }
+}
diff --git a/modules/elastic/common/config.go b/modules/elastic/common/config.go
index eeadb3dd8..4c27a29d4 100644
--- a/modules/elastic/common/config.go
+++ b/modules/elastic/common/config.go
@@ -87,7 +87,9 @@ func InitClientWithConfig(esConfig elastic.ElasticsearchConfig) (client elastic.
ver string
)
if esConfig.Version == "" || esConfig.Version == "auto" {
- verInfo, err := adapter.ClusterVersion(elastic.GetOrInitMetadata(&esConfig))
+ probeMeta := &elastic.ElasticsearchMetadata{Config: &esConfig}
+ probeMeta.Init(true)
+ verInfo, err := adapter.ClusterVersion(probeMeta)
if err != nil {
return nil, err
}
@@ -219,6 +221,9 @@ func InitElasticInstance(esConfig elastic.ElasticsearchConfig) (elastic.API, err
log.Warn("elasticsearch ", esConfig.Name, " is not enabled")
return nil, nil
}
+ originMeta := elastic.GetMetadata(esConfig.ID)
+ initHealth := getInitialMetadataHealth(originMeta)
+
client, err := InitClientWithConfig(esConfig)
if err != nil {
log.Error("elasticsearch ", esConfig.Name, err)
@@ -226,12 +231,6 @@ func InitElasticInstance(esConfig elastic.ElasticsearchConfig) (elastic.API, err
}
elastic.RegisterInstance(esConfig, client)
- originMeta := elastic.GetMetadata(esConfig.ID)
- initHealth := true
- if originMeta != nil {
- initHealth = originMeta.IsAvailable()
- }
-
v := elastic.InitMetadata(&esConfig, initHealth)
if v.Health == nil && originMeta != nil {
v.Health = originMeta.Health
@@ -240,6 +239,13 @@ func InitElasticInstance(esConfig elastic.ElasticsearchConfig) (elastic.API, err
return client, err
}
+func getInitialMetadataHealth(originMeta *elastic.ElasticsearchMetadata) bool {
+ if originMeta == nil {
+ return true
+ }
+ return originMeta.IsAvailable()
+}
+
func GetBasicAuth(esConfig *elastic.ElasticsearchConfig) (basicAuth *model.BasicAuth, err error) {
if esConfig.BasicAuth != nil && esConfig.BasicAuth.Username != "" {
basicAuth = esConfig.BasicAuth
diff --git a/modules/elastic/common/config_test.go b/modules/elastic/common/config_test.go
new file mode 100644
index 000000000..461d5286d
--- /dev/null
+++ b/modules/elastic/common/config_test.go
@@ -0,0 +1,27 @@
+package common
+
+import (
+ "testing"
+
+ "infini.sh/framework/core/elastic"
+)
+
+func TestGetInitialMetadataHealthDefaultsToAvailableForNewCluster(t *testing.T) {
+ if !getInitialMetadataHealth(nil) {
+ t.Fatal("expected new cluster metadata to start as available before first health check")
+ }
+}
+
+func TestGetInitialMetadataHealthKeepsExistingAvailability(t *testing.T) {
+ meta := &elastic.ElasticsearchMetadata{Config: &elastic.ElasticsearchConfig{Enabled: true}}
+ meta.Init(false)
+
+ if getInitialMetadataHealth(meta) {
+ t.Fatal("expected existing unavailable metadata to remain unavailable")
+ }
+
+ meta.Init(true)
+ if !getInitialMetadataHealth(meta) {
+ t.Fatal("expected existing available metadata to remain available")
+ }
+}
diff --git a/modules/elastic/metadata.go b/modules/elastic/metadata.go
index 60c07bdd9..702edc39b 100644
--- a/modules/elastic/metadata.go
+++ b/modules/elastic/metadata.go
@@ -48,49 +48,87 @@ import (
"infini.sh/framework/core/util"
)
+const elasticMetadataKVRetention = 30 * 24 * time.Hour
+
+func shouldRegisterDiscoveredHostForAvailability(meta *elastic.ElasticsearchMetadata, host string) bool {
+ host = util.UnifyLocalAddress(strings.TrimSpace(host))
+ if host == "" {
+ return false
+ }
+
+ if meta == nil {
+ return true
+ }
+ if meta.Config == nil {
+ return true
+ }
+ if meta.Config.Host == "" && len(meta.Config.Hosts) == 0 && meta.Config.Endpoint == "" && len(meta.Config.Endpoints) == 0 {
+ return true
+ }
+
+ seedHosts := meta.GetSeedHosts()
+ if len(seedHosts) == 0 {
+ return true
+ }
+
+ for _, seedHost := range seedHosts {
+ if util.UnifyLocalAddress(strings.TrimSpace(seedHost)) == host {
+ return true
+ }
+ }
+
+ return false
+}
+
func (module *ElasticModule) clusterHealthCheck(clusterID string, force bool) {
log.Tracef("execute health check for: %v", clusterID)
cfg := elastic.GetConfig(clusterID)
+ if cfg == nil || !cfg.Enabled {
+ return
+ }
+
+ if !force && !cfg.Monitored {
+ log.Tracef("skip health check for unmonitored cluster: %v", clusterID)
+ return
+ }
metadata := elastic.GetOrInitMetadata(cfg)
- if cfg.Enabled || force {
- //check seeds' availability
- if force {
- //add seeds to host for health check
- hosts := metadata.GetSeedHosts()
- for _, host := range hosts {
- elastic.GetOrInitHost(host, clusterID)
- }
+ //check seeds' availability
+ if force {
+ //add seeds to host for health check
+ hosts := metadata.GetSeedHosts()
+ for _, host := range hosts {
+ elastic.GetOrInitHost(host, clusterID)
}
- //metadata.GetHttpClient(metadata.GetActivePreferredSeedEndpoint())
- client := elastic.GetClient(cfg.ID)
- //check cluster health status
- health, err := client.ClusterHealth(nil)
- if err != nil || health == nil || health.StatusCode != 200 {
- if health != nil && util.ContainStr(util.UnsafeBytesToString(health.RawResult.Body), "master_not_discovered_exception") {
- metadata.ReportFailure(errors.New("master_not_discovered_exception"))
- } else {
- metadata.ReportFailure(err)
- }
- if metadata.Config.Source == elastic.ElasticsearchConfigSourceElasticsearch && !metadata.IsAvailable() {
- updateClusterHealthStatus(clusterID, "unavailable")
- }
+ }
+ //metadata.GetHttpClient(metadata.GetActivePreferredSeedEndpoint())
+ client := elastic.GetClient(cfg.ID)
+ //check cluster health status
+ health, err := client.ClusterHealth(nil)
+ if err != nil || health == nil || health.StatusCode != 200 {
+ if health != nil && util.ContainStr(util.UnsafeBytesToString(health.RawResult.Body), "master_not_discovered_exception") {
+ metadata.ReportFailure(errors.New("master_not_discovered_exception"))
} else {
- if metadata.Health == nil || metadata.Health.NumberOfNodes == 0 || metadata.Health.Status != health.Status || !metadata.IsAvailable() || force {
- if metadata.Config.Source == elastic.ElasticsearchConfigSourceElasticsearch {
- updateClusterHealthStatus(clusterID, health.Status)
- }
- log.Tracef("cluster [%v] health [%v] updated", clusterID, metadata.Health)
- }
- changes, err := util.DiffTwoObject(metadata.Health, health)
- if err != nil {
- log.Errorf("diff cluster health error: %v", err)
- }
- metadata.ReportSuccess()
- if len(changes) > 0 {
- metadata.Health = health
+ metadata.ReportFailure(err)
+ }
+ if metadata.Config.Source == elastic.ElasticsearchConfigSourceElasticsearch && !metadata.IsAvailable() {
+ updateClusterHealthStatus(clusterID, "unavailable")
+ }
+ } else {
+ if metadata.Health == nil || metadata.Health.NumberOfNodes == 0 || metadata.Health.Status != health.Status || !metadata.IsAvailable() || force {
+ if metadata.Config.Source == elastic.ElasticsearchConfigSourceElasticsearch {
+ updateClusterHealthStatus(clusterID, health.Status)
}
+ log.Tracef("cluster [%v] health [%v] updated", clusterID, metadata.Health)
+ }
+ changes, err := util.DiffTwoObject(metadata.Health, health)
+ if err != nil {
+ log.Errorf("diff cluster health error: %v", err)
+ }
+ metadata.ReportSuccess()
+ if len(changes) > 0 {
+ metadata.Health = health
}
}
}
@@ -177,18 +215,63 @@ func updateClusterHealthStatus(clusterID string, healthStatus string) {
}
+func SyncClusterHealthStatus(clusterID string) {
+ if strings.TrimSpace(clusterID) == "" {
+ return
+ }
+
+ metadata := elastic.GetMetadata(clusterID)
+ if metadata == nil || metadata.Config == nil {
+ return
+ }
+ if metadata.Config.Source != elastic.ElasticsearchConfigSourceElasticsearch {
+ return
+ }
+
+ healthStatus := "unavailable"
+ if metadata.IsAvailable() {
+ if client := elastic.GetClientNoPanic(clusterID); client != nil {
+ health, err := client.ClusterHealth(nil)
+ if err == nil && health != nil && health.StatusCode == 200 && strings.TrimSpace(health.Status) != "" {
+ metadata.Health = health
+ healthStatus = health.Status
+ } else if metadata.Health != nil && strings.TrimSpace(metadata.Health.Status) != "" {
+ healthStatus = metadata.Health.Status
+ } else {
+ healthStatus = "green"
+ }
+ } else if metadata.Health != nil && strings.TrimSpace(metadata.Health.Status) != "" {
+ healthStatus = metadata.Health.Status
+ } else {
+ healthStatus = "green"
+ }
+ }
+
+ updateClusterHealthStatus(clusterID, healthStatus)
+}
+
// update cluster state, on state version change
func (module *ElasticModule) updateClusterState(clusterId string, force bool) {
+ startAt := time.Now()
meta := elastic.GetMetadata(clusterId)
if meta == nil {
return
}
+ if !force && !meta.Config.Monitored {
+ return
+ }
if !force && !meta.IsAvailable() {
return
}
+ interval := moduleConfig.MetadataRefresh.Interval
+ if meta.Config != nil && meta.Config.MetadataConfigs != nil && meta.Config.MetadataConfigs.MetadataRefresh.Interval != "" {
+ interval = meta.Config.MetadataConfigs.MetadataRefresh.Interval
+ }
+ intervalD := util.GetDurationOrDefault(interval, 30*time.Second)
+
client := elastic.GetClient(clusterId)
state, err := client.GetClusterState()
if err != nil {
@@ -201,6 +284,18 @@ func (module *ElasticModule) updateClusterState(clusterId string, force bool) {
}
if state != nil {
+ responseSize := uint64(0)
+ if state.RawResult != nil {
+ responseSize = state.RawResult.Size
+ }
+ indexCount := 0
+ if state.Metadata != nil {
+ indexCount = len(state.Metadata.Indices)
+ }
+ routingIndexCount := 0
+ if state.RoutingTable != nil {
+ routingIndexCount = len(state.RoutingTable.Indices)
+ }
stateChanged := false
if meta.ClusterState == nil {
stateChanged = true
@@ -210,25 +305,23 @@ func (module *ElasticModule) updateClusterState(clusterId string, force bool) {
log.Tracef("cluster state updated from version [%v] to [%v]", meta.ClusterState.Version, state.Version)
}
- oldIndexState, err := kv.GetCompressedValue(elastic.KVElasticIndexMetadata, []byte(clusterId))
-
- //TODO locker
- if stateChanged || (err == nil && oldIndexState == nil) {
- if meta.Config.Source == elastic.ElasticsearchConfigSourceElasticsearch {
- if meta.ClusterState == nil || oldIndexState == nil {
- //load init state from es when console start
- oldIndexState, err = module.loadIndexMetadataFromES(clusterId)
- if err != nil {
- log.Errorf("failed to load index metadata from es: %v", err)
- }
- err = kv.AddValueCompress(elastic.KVElasticIndexMetadata, []byte(clusterId), oldIndexState)
- if err != nil {
+ oldIndexState, oldIndexStateErr := kv.GetCompressedValue(elastic.KVElasticIndexMetadata, []byte(clusterId))
+ if meta.Config.Source == elastic.ElasticsearchConfigSourceElasticsearch {
+ if meta.ClusterState == nil || oldIndexState == nil {
+ // load init state from es when console start
+ oldIndexState, oldIndexStateErr = module.loadIndexMetadataFromES(clusterId)
+ if oldIndexStateErr != nil {
+ log.Errorf("failed to load index metadata from es: %v", oldIndexStateErr)
+ } else {
+ if err := kv.AddValueCompressWithTTL(elastic.KVElasticIndexMetadata, []byte(clusterId), oldIndexState, elasticMetadataKVRetention); err != nil {
log.Errorf("failed to save index metadata: %v", err)
}
}
- if err == nil {
- module.saveIndexMetadata(state, clusterId)
- }
+ }
+ // Always run metadata sync in refresh loop. Some distributions may not
+ // bump cluster state version for every index-state change.
+ if oldIndexStateErr == nil {
+ module.saveIndexMetadata(state, clusterId)
}
}
if stateChanged {
@@ -240,6 +333,20 @@ func (module *ElasticModule) updateClusterState(clusterId string, force bool) {
state.Metadata = metaData
meta.ClusterState = state
}
+ elapsed := time.Since(startAt)
+ if elapsed > intervalD {
+ log.Warnf(
+ "refresh cluster state for cluster [%s] completed slowly, elapsed: %v, interval: %s, response_size: %d bytes, compressed_size_in_bytes: %d, metadata_indices: %d, routing_indices: %d, state_version: %d",
+ meta.Config.Name,
+ elapsed,
+ interval,
+ responseSize,
+ state.CompressedSizeInBytes,
+ indexCount,
+ routingIndexCount,
+ state.Version,
+ )
+ }
}
}
@@ -684,7 +791,7 @@ func (module *ElasticModule) saveIndexMetadata(state *elastic.ClusterState, clus
}
if isIndicesStateChange {
- err = kv.AddValueCompress(elastic.KVElasticIndexMetadata, []byte(clusterID), util.MustToJSONBytes(newIndexMetadata))
+ err = kv.AddValueCompressWithTTL(elastic.KVElasticIndexMetadata, []byte(clusterID), util.MustToJSONBytes(newIndexMetadata), elasticMetadataKVRetention)
if err != nil {
log.Error(err)
}
@@ -696,11 +803,18 @@ func (module *ElasticModule) updateNodeInfo(meta *elastic.ElasticsearchMetadata,
log.Trace("update node info")
+ if !force && !meta.Config.Monitored {
+ return
+ }
+
if !force && !meta.IsAvailable() {
+ stateChanged := false
if !force {
- setNodeUnknown(meta.Config.ID)
+ stateChanged = setNodeUnknown(meta.Config.ID)
+ }
+ if stateChanged || rate.GetRateLimiter("metadata_node_info_skip", meta.Config.ID, 1, 1, 10*time.Minute).Allow() {
+ log.Debugf("elasticsearch [%v] is not available, skip update node info", meta.Config.Name)
}
- log.Debugf("elasticsearch [%v] is not available, skip update node info", meta.Config.Name)
return
}
@@ -767,7 +881,11 @@ func (module *ElasticModule) updateNodeInfo(meta *elastic.ElasticsearchMetadata,
if moduleConfig.ORMConfig.Enabled {
if meta.Config.Source == elastic.ElasticsearchConfigSourceElasticsearch {
//todo check whether store elasticsearch change or not
- err = saveNodeMetadata(*nodes, meta.Config.ID)
+ clusterUUID := meta.Config.ClusterUUID
+ if meta.ClusterState != nil && meta.ClusterState.ClusterUUID != "" {
+ clusterUUID = meta.ClusterState.ClusterUUID
+ }
+ err = saveNodeMetadata(*nodes, meta.Config.ID, clusterUUID)
if err != nil {
if rate.GetRateLimiterPerSecond(meta.Config.ID, "save_nodes_metadata_on_error", 1).Allow() {
log.Errorf("elasticsearch [%v] failed to save nodes info: %v", meta.Config.Name, err)
@@ -786,7 +904,11 @@ func (module *ElasticModule) updateNodeInfo(meta *elastic.ElasticsearchMetadata,
//register host to do availability monitoring
if discovery {
for _, v := range *nodes {
- elastic.GetOrInitHost(v.GetHttpPublishHost(), meta.Config.ID)
+ host := v.GetHttpPublishHost()
+ if !shouldRegisterDiscoveredHostForAvailability(meta, host) {
+ continue
+ }
+ elastic.GetOrInitHost(host, meta.Config.ID)
}
}
@@ -795,7 +917,7 @@ func (module *ElasticModule) updateNodeInfo(meta *elastic.ElasticsearchMetadata,
"nodes": nodes,
"timestamp": time.Now(),
}
- err = kv.AddValueCompress(elastic.KVElasticNodeMetadata, []byte(meta.Config.ID), util.MustToJSONBytes(cacheNodeInfo))
+ err = kv.AddValueCompressWithTTL(elastic.KVElasticNodeMetadata, []byte(meta.Config.ID), util.MustToJSONBytes(cacheNodeInfo), elasticMetadataKVRetention)
if err != nil {
log.Errorf("save node metadata error: %v", err)
}
@@ -808,17 +930,17 @@ func (module *ElasticModule) updateNodeInfo(meta *elastic.ElasticsearchMetadata,
var saveNodeMetadataMutex = sync.Mutex{}
var nodeAlreadyUnknown = map[string]bool{}
-func setNodeUnknown(clusterID string) {
+func setNodeUnknown(clusterID string) bool {
kv.DeleteKey(elastic.KVElasticNodeMetadata, []byte(clusterID))
meta := elastic.GetMetadata(clusterID)
if meta == nil {
- return
+ return false
}
if meta.Config.Source != elastic.ElasticsearchConfigSourceElasticsearch {
- return
+ return false
}
if v, ok := nodeAlreadyUnknown[clusterID]; ok && v {
- return
+ return false
}
queueConfig := queue.GetOrInitConfig(elastic.QueueElasticIndexState)
if queueConfig.Labels == nil {
@@ -846,8 +968,9 @@ func setNodeUnknown(clusterID string) {
}
nodeAlreadyUnknown[clusterID] = true
+ return true
}
-func saveNodeMetadata(nodes map[string]elastic.NodesInfo, clusterID string) error {
+func saveNodeMetadata(nodes map[string]elastic.NodesInfo, clusterID, clusterUUID string) error {
esConfig := elastic.GetConfig(clusterID)
saveNodeMetadataMutex.Lock()
defer func() {
@@ -857,28 +980,63 @@ func saveNodeMetadata(nodes map[string]elastic.NodesInfo, clusterID string) erro
}
}()
- queryDslTpl := `{
- "size": 1000,
- "query": {
- "bool": {
- "must": [
- {"term": {
- "metadata.cluster_id": {
- "value": "%s"
- }
- }},
- {"term": {
- "metadata.category": {
- "value": "elasticsearch"
- }
- }}
- ]
- }
- }
-}`
- queryDsl := fmt.Sprintf(queryDslTpl, clusterID)
+ logicalClusterID := clusterID
+ if clusterUUID != "" {
+ logicalClusterID = clusterUUID
+ }
+ nodeDocID := func(clusterKey, nodeID string) string {
+ return util.MD5digest(fmt.Sprintf("%s:%s", clusterKey, nodeID))
+ }
+
+ must := []util.MapStr{
+ {
+ "term": util.MapStr{
+ "metadata.category": util.MapStr{
+ "value": "elasticsearch",
+ },
+ },
+ },
+ }
+ boolQuery := util.MapStr{
+ "must": must,
+ }
+ if clusterUUID != "" {
+ boolQuery["should"] = []util.MapStr{
+ {
+ "term": util.MapStr{
+ "metadata.cluster_id": util.MapStr{
+ "value": clusterID,
+ },
+ },
+ },
+ {
+ "term": util.MapStr{
+ "metadata.labels.cluster_uuid": util.MapStr{
+ "value": clusterUUID,
+ },
+ },
+ },
+ }
+ boolQuery["minimum_should_match"] = 1
+ } else {
+ must = append(must, util.MapStr{
+ "term": util.MapStr{
+ "metadata.cluster_id": util.MapStr{
+ "value": clusterID,
+ },
+ },
+ })
+ boolQuery["must"] = must
+ }
+
+ queryDsl := util.MustToJSONBytes(util.MapStr{
+ "size": 1000,
+ "query": util.MapStr{
+ "bool": boolQuery,
+ },
+ })
q := &orm.Query{}
- q.RawQuery = []byte(queryDsl)
+ q.RawQuery = queryDsl
err, result := orm.Search(&elastic.NodeConfig{}, q)
if err != nil {
return err
@@ -893,7 +1051,19 @@ func saveNodeMetadata(nodes map[string]elastic.NodesInfo, clusterID string) erro
//nodeMetadatas[nodeID] = nodeInfo
if nid, ok := nodeID.(string); ok {
if id, ok := nodeInfo["id"]; ok {
- nodeIDMap[nid] = id
+ existingID, hasExisting := nodeIDMap[nid]
+ canonicalID := nodeDocID(logicalClusterID, nid)
+ if !hasExisting {
+ nodeIDMap[nid] = id
+ } else {
+ existingIDStr, existingOK := existingID.(string)
+ idStr, currentOK := id.(string)
+ if currentOK && idStr == canonicalID {
+ nodeIDMap[nid] = id
+ } else if !(existingOK && existingIDStr == canonicalID) {
+ nodeIDMap[nid] = id
+ }
+ }
}
historyNodeMetadata[nid] = nodeInfo
if _, ok = nodes[nid]; !ok {
@@ -908,15 +1078,26 @@ func saveNodeMetadata(nodes map[string]elastic.NodesInfo, clusterID string) erro
rawBytes := util.MustToJSONBytes(nodeInfo)
currentNodeInfo := util.MapStr{}
util.MustFromJSONBytes(rawBytes, ¤tNodeInfo)
+ canonicalID := nodeDocID(logicalClusterID, rawNodeID)
+ legacyID := nodeDocID(clusterID, rawNodeID)
var innerID interface{}
var typ string
var changeLog diff.Changelog
if rowID, ok := nodeIDMap[rawNodeID]; !ok {
//new
- newID := fmt.Sprintf("%s:%s", clusterID, rawNodeID)
- newID = util.MD5digest(newID)
+ newID := canonicalID
typ = "create"
innerID = newID
+ labels := util.MapStr{
+ "transport_address": nodeInfo.TransportAddress,
+ "ip": nodeInfo.Ip,
+ "version": nodeInfo.Version,
+ "roles": nodeInfo.Roles,
+ "status": "available",
+ }
+ if clusterUUID != "" {
+ labels["cluster_uuid"] = clusterUUID
+ }
nodeMetadata := &elastic.NodeConfig{
Metadata: elastic.NodeMetadata{
ClusterID: clusterID,
@@ -925,13 +1106,7 @@ func saveNodeMetadata(nodes map[string]elastic.NodesInfo, clusterID string) erro
ClusterName: esConfig.Name,
NodeName: nodeInfo.Name,
Host: nodeInfo.Host,
- Labels: util.MapStr{
- "transport_address": nodeInfo.TransportAddress,
- "ip": nodeInfo.Ip,
- "version": nodeInfo.Version,
- "roles": nodeInfo.Roles,
- "status": "available",
- },
+ Labels: labels,
},
ID: newID,
Timestamp: time.Now(),
@@ -943,7 +1118,7 @@ func saveNodeMetadata(nodes map[string]elastic.NodesInfo, clusterID string) erro
log.Error(err)
}
} else {
- innerID = rowID
+ innerID = canonicalID
typ = "update"
if rid, ok := rowID.(string); ok {
if historyM, ok := historyNodeMetadata[rawNodeID]; ok {
@@ -964,6 +1139,9 @@ func saveNodeMetadata(nodes map[string]elastic.NodesInfo, clusterID string) erro
"roles": nodeInfo.Roles,
"status": "available",
}
+ if clusterUUID != "" {
+ newLabels["cluster_uuid"] = clusterUUID
+ }
if labels, err := historyM.GetValue("metadata.labels"); err == nil {
if labelsM, ok := labels.(map[string]interface{}); ok {
if st, ok := labelsM["status"].(string); ok && st == "unavailable" || st == "N/A" {
@@ -1008,7 +1186,7 @@ func saveNodeMetadata(nodes map[string]elastic.NodesInfo, clusterID string) erro
Labels: newLabels,
Category: "elasticsearch",
},
- ID: rid,
+ ID: canonicalID,
Timestamp: time.Now(),
Payload: elastic.NodePayload{NodeInfo: &nodeInfo},
}
@@ -1017,6 +1195,14 @@ func saveNodeMetadata(nodes map[string]elastic.NodesInfo, clusterID string) erro
if err != nil {
log.Error(err)
}
+ if clusterUUID != "" && legacyID != canonicalID && rid == legacyID {
+ delCtx := orm.NewContext().DirectAccess()
+ delCtx.Set(orm.CheckExistsBeforeDelete, false)
+ err = orm.Delete(delCtx, &elastic.NodeConfig{ID: legacyID})
+ if err != nil {
+ log.Error(err)
+ }
+ }
}
}
@@ -1065,6 +1251,12 @@ func saveNodeMetadata(nodes map[string]elastic.NodesInfo, clusterID string) erro
if oldStatus, ok := oldConfig.Metadata.Labels["status"].(string); ok && oldStatus == "unavailable" {
continue
}
+ if oldConfig.Metadata.Labels == nil {
+ oldConfig.Metadata.Labels = util.MapStr{}
+ }
+ if clusterUUID != "" {
+ oldConfig.Metadata.Labels["cluster_uuid"] = clusterUUID
+ }
oldConfig.Metadata.Labels["status"] = "unavailable"
oldConfig.Timestamp = time.Now()
@@ -1144,6 +1336,9 @@ func saveNodeMetadata(nodes map[string]elastic.NodesInfo, clusterID string) erro
// on demand, on state version change
func updateAliases(meta *elastic.ElasticsearchMetadata, force bool) {
+ if !force && !meta.Config.Monitored {
+ return
+ }
if !force && !meta.IsAvailable() {
return
@@ -1244,6 +1439,9 @@ func (module *ElasticModule) updateClusterSettings(clusterId string) {
if meta == nil {
return
}
+ if !meta.Config.Monitored {
+ return
+ }
if !meta.IsAvailable() {
return
}
@@ -1318,9 +1516,15 @@ func (module *ElasticModule) updateClusterSettings(clusterId string) {
if err != nil {
panic(err)
}
- kv.AddValue(elastic.KVElasticClusterSettings, []byte(clusterId), util.MustToJSONBytes(settings))
+ err = kv.AddValueWithTTL(elastic.KVElasticClusterSettings, []byte(clusterId), util.MustToJSONBytes(settings), elasticMetadataKVRetention)
+ if err != nil {
+ log.Errorf("failed to save cluster settings: %v", err)
+ }
} else {
- kv.AddValue(elastic.KVElasticClusterSettings, []byte(clusterId), util.MustToJSONBytes(settings))
+ err = kv.AddValueWithTTL(elastic.KVElasticClusterSettings, []byte(clusterId), util.MustToJSONBytes(settings), elasticMetadataKVRetention)
+ if err != nil {
+ log.Errorf("failed to save cluster settings: %v", err)
+ }
}
}
diff --git a/modules/elastic/metadata_discovery_test.go b/modules/elastic/metadata_discovery_test.go
new file mode 100644
index 000000000..4db2b96b6
--- /dev/null
+++ b/modules/elastic/metadata_discovery_test.go
@@ -0,0 +1,38 @@
+package elastic
+
+import (
+ "testing"
+
+ coreelastic "infini.sh/framework/core/elastic"
+ "infini.sh/framework/core/orm"
+)
+
+func TestShouldRegisterDiscoveredHostForAvailabilityPrefersSeedHosts(t *testing.T) {
+ meta := &coreelastic.ElasticsearchMetadata{
+ Config: &coreelastic.ElasticsearchConfig{
+ ORMObjectBase: orm.ORMObjectBase{ID: "cluster-1"},
+ Host: "192.168.3.8:9200",
+ Hosts: []string{"192.168.3.8:9200"},
+ },
+ }
+
+ if shouldRegisterDiscoveredHostForAvailability(meta, "172.22.0.2:9200") {
+ t.Fatal("expected non-seed discovered host to be excluded from availability monitoring")
+ }
+
+ if !shouldRegisterDiscoveredHostForAvailability(meta, "192.168.3.8:9200") {
+ t.Fatal("expected seed host to remain eligible for availability monitoring")
+ }
+}
+
+func TestShouldRegisterDiscoveredHostForAvailabilityAllowsDiscoveryWithoutSeeds(t *testing.T) {
+ meta := &coreelastic.ElasticsearchMetadata{
+ Config: &coreelastic.ElasticsearchConfig{
+ ORMObjectBase: orm.ORMObjectBase{ID: "cluster-2"},
+ },
+ }
+
+ if !shouldRegisterDiscoveredHostForAvailability(meta, "172.22.0.2:9200") {
+ t.Fatal("expected discovered host to be eligible when no seed hosts are configured")
+ }
+}
diff --git a/modules/elastic/module.go b/modules/elastic/module.go
index 2e26215ce..59f3cbd5d 100755
--- a/modules/elastic/module.go
+++ b/modules/elastic/module.go
@@ -113,10 +113,26 @@ func loadFileBasedElasticConfig() []elastic.ElasticsearchConfig {
return configs
}
+func lookupSystemElasticsearchID() (string, bool) {
+ value := global.Lookup(elastic.GlobalSystemElasticsearchID)
+ systemID, ok := value.(string)
+ if !ok || systemID == "" {
+ return "", false
+ }
+ return systemID, true
+}
+
func loadESBasedElasticConfig() []elastic.ElasticsearchConfig {
configs := []elastic.ElasticsearchConfig{}
+ systemID, ok := lookupSystemElasticsearchID()
+ if !ok {
+ return configs
+ }
query := elastic.SearchRequest{From: 0, Size: 1000} //TODO handle clusters beyond 1000
- esClient := elastic.GetClient(global.MustLookupString(elastic.GlobalSystemElasticsearchID))
+ query.Set("query", util.MapStr{
+ "match_all": util.MapStr{},
+ })
+ esClient := elastic.GetClient(systemID)
result, err := esClient.Search(orm.GetIndexName(elastic.ElasticsearchConfig{}), &query)
if err != nil {
log.Error(err)
@@ -246,7 +262,7 @@ func nodeAvailabilityCheck() {
}
cfg := elastic.GetConfig(v.ClusterID)
- if !cfg.Enabled || (cfg.MetadataConfigs != nil && !cfg.MetadataConfigs.NodeAvailabilityCheck.Enabled) {
+ if !cfg.Enabled || !cfg.Monitored || (cfg.MetadataConfigs != nil && !cfg.MetadataConfigs.NodeAvailabilityCheck.Enabled) {
return true
}
@@ -265,6 +281,7 @@ func nodeAvailabilityCheck() {
}
if time.Since(startTime.(time.Time)) > util.GetDurationOrDefault(interval, 10*time.Second)*2 {
log.Warnf("check availability for node [%s] is still running, elapsed: %v, skip waiting", v.Host, elapsed.String())
+ return true
} else {
log.Warnf("check availability for node [%s] is still running, elapsed: %v", v.Host, elapsed.String())
return true
@@ -313,7 +330,7 @@ func (module *ElasticModule) registerClusterStateRefreshTask() {
log.Tracef("init meta refresh task: [%v] [%v] [%v] [%v]", key, v.ID, v.Name, v.Enabled)
if ok {
- if !v.Enabled || (v.MetadataConfigs != nil && !v.MetadataConfigs.MetadataRefresh.Enabled) {
+ if !v.Enabled || !v.Monitored || (v.MetadataConfigs != nil && !v.MetadataConfigs.MetadataRefresh.Enabled) {
return true
}
@@ -326,6 +343,7 @@ func (module *ElasticModule) registerClusterStateRefreshTask() {
intervalD := util.GetDurationOrDefault(interval, 10*time.Second)
if time.Since(startTime.(time.Time)) > intervalD*2 {
log.Warnf("refresh cluster state for cluster [%s] is still running, elapsed: %v, skip waiting", v.Name, elapsed.String())
+ return true
} else {
duration := elapsed - intervalD
abd := math.Abs(duration.Seconds())
@@ -339,8 +357,8 @@ func (module *ElasticModule) registerClusterStateRefreshTask() {
task.RunWithContext("refresh_cluster_state", func(ctx context.Context) error {
clusterID := task.MustGetString(ctx, "id")
+ defer module.stateMap.Delete(clusterID)
module.updateClusterState(clusterID, false)
- module.stateMap.Delete(clusterID)
return nil
}, context.WithValue(context.Background(), "id", v.ID))
}
@@ -395,20 +413,37 @@ func InitSchema() {
var ormInited bool
func (module *ElasticModule) Start() error {
+ systemID, hasSystemCluster := lookupSystemElasticsearchID()
if moduleConfig.ORMConfig.Enabled {
- client := elastic.GetClient(global.MustLookupString(elastic.GlobalSystemElasticsearchID))
- handler := ElasticORM{Client: client, Config: moduleConfig.ORMConfig}
- orm.Register("elastic", &handler)
+ if !hasSystemCluster {
+ log.Warn("skip elastic ORM initialization, system cluster is not available")
+ } else {
+ client := elastic.GetClient(systemID)
+ handler := ElasticORM{Client: client, Config: moduleConfig.ORMConfig}
+ if orm.HasAdapter("elastic") {
+ log.Debug("skip duplicate elastic ORM registration")
+ } else {
+ orm.Register("elastic", &handler)
+ }
+ }
}
if moduleConfig.StoreConfig.Enabled {
- client := elastic.GetClient(global.MustLookupString(elastic.GlobalSystemElasticsearchID))
- module.storeHandler = &ElasticStore{Client: client, Config: moduleConfig.StoreConfig}
- kv.Register("elastic", module.storeHandler)
+ if !hasSystemCluster {
+ log.Warn("skip elastic store initialization, system cluster is not available")
+ } else {
+ client := elastic.GetClient(systemID)
+ module.storeHandler = &ElasticStore{Client: client, Config: moduleConfig.StoreConfig}
+ if kv.HasStore("elastic") {
+ log.Debug("skip duplicate elastic store registration")
+ } else {
+ kv.Register("elastic", module.storeHandler)
+ }
+ }
}
- if moduleConfig.ORMConfig.Enabled {
+ if moduleConfig.ORMConfig.Enabled && hasSystemCluster {
if !ormInited {
//init template
InitTemplate(false)
@@ -419,8 +454,12 @@ func (module *ElasticModule) Start() error {
}
if moduleConfig.RemoteConfigEnabled {
- m := loadESBasedElasticConfig()
- initElasticInstances(m, elastic.ElasticsearchConfigSourceElasticsearch)
+ if !hasSystemCluster {
+ log.Warn("skip remote elastic config loading, system cluster is not available")
+ } else {
+ m := loadESBasedElasticConfig()
+ initElasticInstances(m, elastic.ElasticsearchConfigSourceElasticsearch)
+ }
}
if module.storeHandler != nil {
@@ -442,23 +481,19 @@ func (module *ElasticModule) Start() error {
cfg1, ok := value.(*elastic.ElasticsearchConfig)
if ok && cfg1 != nil {
log.Tracef("init elasticsearch config: %v", cfg1.Name)
- metadata := elastic.GetMetadata(cfg1.ID)
- if metadata != nil {
- //update nodes
- module.updateNodeInfo(metadata, true, cfg1.Discovery.Enabled)
+ if cfg1.Monitored {
+ metadata := elastic.GetMetadata(cfg1.ID)
+ if metadata != nil {
+ //update nodes
+ module.updateNodeInfo(metadata, true, cfg1.Discovery.Enabled)
- //update alias
- updateAliases(metadata, true)
+ //update alias
+ updateAliases(metadata, true)
- //update
- module.updateClusterState(cfg1.ID, true)
+ //update
+ module.updateClusterState(cfg1.ID, true)
+ }
}
-
- task.RunWithContext("cluster_health_check", func(ctx context.Context) error {
- id := task.MustGetString(ctx, "id")
- module.clusterHealthCheck(id, true)
- return nil
- }, context.WithValue(context.Background(), "id", cfg1.ID))
}
return true
})
@@ -476,7 +511,7 @@ func (module *ElasticModule) Start() error {
}
cfg1, ok := value.(*elastic.ElasticsearchConfig)
if ok && cfg1 != nil {
- if !cfg1.Enabled || (cfg1.MetadataConfigs != nil && !cfg1.MetadataConfigs.HealthCheck.Enabled) {
+ if !cfg1.Enabled || !cfg1.Monitored || (cfg1.MetadataConfigs != nil && !cfg1.MetadataConfigs.HealthCheck.Enabled) {
return true
}
@@ -491,6 +526,7 @@ func (module *ElasticModule) Start() error {
tinterval := util.GetDurationOrDefault(interval, 10*time.Second)
if elapsed > tinterval*2 {
log.Warnf("health check for cluster [%s] is still running, elapsed: %v, skip waiting", cfg1.Name, elapsed.String())
+ return true
} else if math.Abs((elapsed - tinterval).Seconds()) > 3 {
log.Warnf("health check for cluster [%s] is still running, elapsed: %v", cfg1.Name, elapsed.String())
return true
@@ -500,8 +536,8 @@ func (module *ElasticModule) Start() error {
task.RunWithContext("refresh_cluster_health", func(ctx context.Context) error {
clusterID := task.MustGetString(ctx, "id")
+ defer module.healthMap.Delete(clusterID)
module.clusterHealthCheck(clusterID, false)
- module.healthMap.Delete(clusterID)
return nil
}, context.WithValue(context.Background(), "id", cfg1.ID))
}
@@ -646,7 +682,7 @@ func (module *ElasticModule) registerClusterSettingsRefreshTask() {
log.Tracef("init settings refresh task: [%v] [%v] [%v] [%v]", key, v.ID, v.Name, v.Enabled)
if ok {
- if !v.Enabled || (v.MetadataConfigs != nil && !v.MetadataConfigs.ClusterSettingsCheck.Enabled) {
+ if !v.Enabled || !v.Monitored || (v.MetadataConfigs != nil && !v.MetadataConfigs.ClusterSettingsCheck.Enabled) {
return true
}
if startTime, ok := module.settingsMap.Load(v.ID); ok {
@@ -658,6 +694,7 @@ func (module *ElasticModule) registerClusterSettingsRefreshTask() {
if time.Since(startTime.(time.Time)) > util.GetDurationOrDefault(interval, 10*time.Second)*2 {
log.Warnf("collect cluster settings for cluster [%s] is still running, elapsed: %v, skip waiting", v.Name, elapsed.String())
+ return true
} else {
log.Warnf("collect cluster settings for cluster [%s] is still running, elapsed: %v", v.Name, elapsed.String())
return true
@@ -666,8 +703,8 @@ func (module *ElasticModule) registerClusterSettingsRefreshTask() {
module.settingsMap.Store(v.ID, time.Now())
task.RunWithContext("refresh_cluster_settings", func(ctx context.Context) error {
clusterID := task.MustGetString(ctx, "id")
+ defer module.settingsMap.Delete(clusterID)
module.updateClusterSettings(clusterID)
- module.settingsMap.Delete(clusterID)
return nil
}, context.WithValue(context.Background(), "id", v.ID))
}
@@ -694,7 +731,18 @@ func (module *ElasticModule) refreshAllClusterMetadata() {
log.Trace("update elasticsearch's metadata:", v, ok)
if ok {
- module.updateNodeInfo(v, false, v.Config.Discovery.Enabled)
+ cfg := elastic.GetConfigNoPanic(v.Config.ID)
+ if cfg == nil {
+ log.Debugf("elasticsearch metadata [%v] has no active config, removing stale metadata", v.Config.ID)
+ elastic.RemoveInstance(v.Config.ID)
+ elastic.RemoveHostsByClusterID(v.Config.ID)
+ return true
+ }
+ v.Config = cfg
+ if !cfg.Enabled || !cfg.Monitored || (cfg.MetadataConfigs != nil && !cfg.MetadataConfigs.MetadataRefresh.Enabled) {
+ return true
+ }
+ module.updateNodeInfo(v, false, cfg.Discovery.Enabled)
}
return true
})
@@ -707,6 +755,17 @@ func (module *ElasticModule) refreshAllClusterAlias(force bool) {
}
v, ok := value.(*elastic.ElasticsearchMetadata)
if ok {
+ cfg := elastic.GetConfigNoPanic(v.Config.ID)
+ if cfg == nil {
+ log.Debugf("elasticsearch metadata [%v] has no active config, removing stale metadata", v.Config.ID)
+ elastic.RemoveInstance(v.Config.ID)
+ elastic.RemoveHostsByClusterID(v.Config.ID)
+ return true
+ }
+ v.Config = cfg
+ if !cfg.Enabled || !cfg.Monitored || (cfg.MetadataConfigs != nil && !cfg.MetadataConfigs.MetadataRefresh.Enabled) {
+ return true
+ }
updateAliases(v, force)
}
return true
diff --git a/modules/elastic/module_test.go b/modules/elastic/module_test.go
index 3accbf667..e88f7e39c 100644
--- a/modules/elastic/module_test.go
+++ b/modules/elastic/module_test.go
@@ -1,48 +1,48 @@
-// Copyright (C) INFINI Labs & INFINI LIMITED.
-//
-// The INFINI Framework is offered under the GNU Affero General Public License v3.0
-// and as commercial software.
-//
-// For commercial licensing, contact us at:
-// - Website: infinilabs.com
-// - Email: hello@infini.ltd
-//
-// Open Source licensed under AGPL V3:
-// This program is free software: you can redistribute it and/or modify
-// it under the terms of the GNU Affero General Public License as published by
-// the Free Software Foundation, either version 3 of the License, or
-// (at your option) any later version.
-//
-// This program is distributed in the hope that it will be useful,
-// but WITHOUT ANY WARRANTY; without even the implied warranty of
-// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-// GNU Affero General Public License for more details.
-//
-// You should have received a copy of the GNU Affero General Public License
-// along with this program. If not, see .
-
package elastic
import (
- "fmt"
- "github.com/buger/jsonparser"
- "infini.sh/framework/core/util"
"testing"
+
+ coreElastic "infini.sh/framework/core/elastic"
+ "infini.sh/framework/core/global"
)
-func TestV7GetClusterStates(t *testing.T) {
- str := "{ \"_nodes\": { \"total\": 1, \"successful\": 1, \"failed\": 0 }, \"cluster_name\": \"es-v700\", \"cluster_uuid\": \"7NtDffC3RzGChhoOmgySig\", \"timestamp\": 1629611578327, \"status\": \"green\", \"indices\": { \"count\": 0, \"shards\": {}, \"docs\": { \"count\": 0, \"deleted\": 0 }, \"store\": { \"size_in_bytes\": 0 }, \"fielddata\": { \"memory_size_in_bytes\": 0, \"evictions\": 0 }, \"query_cache\": { \"memory_size_in_bytes\": 0, \"total_count\": 0, \"hit_count\": 0, \"miss_count\": 0, \"cache_size\": 0, \"cache_count\": 0, \"evictions\": 0 }, \"completion\": { \"size_in_bytes\": 0 }, \"segments\": { \"count\": 0, \"memory_in_bytes\": 0, \"terms_memory_in_bytes\": 0, \"stored_fields_memory_in_bytes\": 0, \"term_vectors_memory_in_bytes\": 0, \"norms_memory_in_bytes\": 0, \"points_memory_in_bytes\": 0, \"doc_values_memory_in_bytes\": 0, \"index_writer_memory_in_bytes\": 0, \"version_map_memory_in_bytes\": 0, \"fixed_bit_set_memory_in_bytes\": 0, \"max_unsafe_auto_id_timestamp\": -9223372036854776000, \"file_sizes\": {} } }, \"nodes\": { \"count\": { \"total\": 1, \"data\": 1, \"coordinating_only\": 0, \"master\": 1, \"ingest\": 1 }, \"versions\": [ \"7.0.0\" ], \"os\": { \"available_processors\": 24, \"allocated_processors\": 24, \"names\": [ { \"name\": \"Windows 10\", \"count\": 1 } ], \"pretty_names\": [ { \"pretty_name\": \"Windows 10\", \"count\": 1 } ], \"mem\": { \"total_in_bytes\": 137121308672, \"free_in_bytes\": 114813546496, \"used_in_bytes\": 22307762176, \"free_percent\": 84, \"used_percent\": 16 } }, \"process\": { \"cpu\": { \"percent\": 0 }, \"open_file_descriptors\": { \"min\": -1, \"max\": -1, \"avg\": 0 } }, \"jvm\": { \"max_uptime_in_millis\": 2021226, \"versions\": [ { \"version\": \"9.0.1.3\", \"vm_name\": \"OpenJDK 64-Bit Server VM\", \"vm_version\": \"9.0.1.3+11\", \"vm_vendor\": \"Azul Systems, Inc.\", \"bundled_jdk\": false, \"using_bundled_jdk\": null, \"count\": 1 } ], \"mem\": { \"heap_used_in_bytes\": 277003800, \"heap_max_in_bytes\": 1037959168 }, \"threads\": 66 }, \"fs\": { \"total_in_bytes\": 6000527532032, \"free_in_bytes\": 3111816585216, \"available_in_bytes\": 3111816585216 }, \"plugins\": [], \"network_types\": { \"transport_types\": { \"netty4\": 1 }, \"http_types\": { \"netty4\": 1 } }, \"discovery_types\": { \"zen\": 1 } } }"
+func TestLoadESBasedElasticConfigSkipsWhenSystemClusterUnavailable(t *testing.T) {
+ previous := global.Lookup(coreElastic.GlobalSystemElasticsearchID)
+ defer global.Register(coreElastic.GlobalSystemElasticsearchID, previous)
+
+ global.Register(coreElastic.GlobalSystemElasticsearchID, "")
+
+ configs := loadESBasedElasticConfig()
+ if len(configs) != 0 {
+ t.Fatalf("expected no remote configs when system cluster id is unavailable, got %d", len(configs))
+ }
+}
+
+func TestElasticModuleStartSkipsSystemClusterDependentInitBeforeSetup(t *testing.T) {
+ previousSystemID := global.Lookup(coreElastic.GlobalSystemElasticsearchID)
+ defer global.Register(coreElastic.GlobalSystemElasticsearchID, previousSystemID)
+
+ previousModuleConfig := moduleConfig
+ defer func() {
+ moduleConfig = previousModuleConfig
+ }()
+
+ previousOrmInited := ormInited
+ defer func() {
+ ormInited = previousOrmInited
+ }()
+
+ global.Register(coreElastic.GlobalSystemElasticsearchID, "")
+
+ moduleConfig = getDefaultConfig()
+ moduleConfig.ORMConfig.Enabled = true
+ moduleConfig.StoreConfig.Enabled = true
+ moduleConfig.RemoteConfigEnabled = true
+ ormInited = false
- d1, err := jsonparser.GetInt(util.UnsafeStringToBytes(str), "indices", "segments", "max_unsafe_auto_id_timestamp")
- fmt.Println("xv:", d1, err)
- if err != nil {
- d, err := jsonparser.Set(util.UnsafeStringToBytes(str), []byte("-1"), "indices", "segments", "max_unsafe_auto_id_timestamp")
- if err == nil {
- str = util.UnsafeBytesToString(d)
- }
+ module := &ElasticModule{}
+ if err := module.Start(); err != nil {
+ t.Fatalf("expected elastic module start to succeed before setup, got %v", err)
}
- d1, err = jsonparser.GetInt(util.UnsafeStringToBytes(str), "indices", "segments", "max_unsafe_auto_id_timestamp")
- fmt.Println("xv:", d1, err)
- //xv,err:=jsonparser.GetInt([]byte(str),"indices.segments.max_unsafe_auto_id_timestamp")
- //fmt.Println("xv:",xv,err)
}
diff --git a/modules/elastic/orm.go b/modules/elastic/orm.go
index 7da6b0bf6..647d18c60 100755
--- a/modules/elastic/orm.go
+++ b/modules/elastic/orm.go
@@ -46,6 +46,17 @@ type ElasticORM struct {
Config common.ORMConfig
}
+func shouldRetrySearchWithoutCollapse(searchResponse *elastic.SearchResponse, collapseField string) bool {
+ if strings.TrimSpace(collapseField) == "" || searchResponse == nil || searchResponse.RawResult == nil {
+ return false
+ }
+ if searchResponse.RawResult.StatusCode != http.StatusBadRequest {
+ return false
+ }
+ body := string(searchResponse.RawResult.Body)
+ return strings.Contains(body, collapseField) && strings.Contains(body, "in order to collapse on")
+}
+
var templateInited bool
func InitTemplate(force bool) {
@@ -548,10 +559,14 @@ func (handler *ElasticORM) Search(t interface{}, q *api.Query) (error, api.Resul
}
if global.Env().IsDebug {
- log.Info(util.MustToJSON(request))
+ log.Trace(util.MustToJSON(request))
}
searchResponse, err = handler.Client.Search(indexName, &request)
+ if err == nil && shouldRetrySearchWithoutCollapse(searchResponse, q.CollapseField) {
+ request.Collapse = nil
+ searchResponse, err = handler.Client.Search(indexName, &request)
+ }
}
if err != nil {
@@ -621,36 +636,38 @@ func (handler *ElasticORM) SearchWithResultItemMapper(resultArray interface{}, i
searchResponse, err = handler.Client.SearchByTemplate(indexName, q.TemplatedQuery.TemplateID, q.TemplatedQuery.Parameters)
} else {
- request.Query = &elastic.Query{}
- boolQuery := elastic.BoolQuery{}
-
- if q.Conds != nil && len(q.Conds) > 0 {
- for _, cond := range q.Conds {
- query := getQuery(cond)
- switch cond.BoolType {
- case api.Filter:
- boolQuery.Filter = append(boolQuery.Filter, query)
- case api.Must:
- boolQuery.Must = append(boolQuery.Must, query)
- case api.MustNot:
- boolQuery.MustNot = append(boolQuery.MustNot, query)
- case api.Should:
- boolQuery.Should = append(boolQuery.Should, query)
+ if q.Filter != nil || q.Conds != nil && len(q.Conds) > 0 {
+ request.Query = &elastic.Query{}
+ boolQuery := elastic.BoolQuery{}
+
+ if q.Conds != nil && len(q.Conds) > 0 {
+ for _, cond := range q.Conds {
+ query := getQuery(cond)
+ switch cond.BoolType {
+ case api.Filter:
+ boolQuery.Filter = append(boolQuery.Filter, query)
+ case api.Must:
+ boolQuery.Must = append(boolQuery.Must, query)
+ case api.MustNot:
+ boolQuery.MustNot = append(boolQuery.MustNot, query)
+ case api.Should:
+ boolQuery.Should = append(boolQuery.Should, query)
+ }
}
}
- }
- if q.Filter != nil {
- filter := getQuery(q.Filter)
- //temp fix for must_not filters
- if q.Filter.BoolType == api.MustNot {
- boolQuery.MustNot = append(boolQuery.MustNot, filter)
- } else {
- boolQuery.Filter = append(boolQuery.Filter, filter)
+ if q.Filter != nil {
+ filter := getQuery(q.Filter)
+ //temp fix for must_not filters
+ if q.Filter.BoolType == api.MustNot {
+ boolQuery.MustNot = append(boolQuery.MustNot, filter)
+ } else {
+ boolQuery.Filter = append(boolQuery.Filter, filter)
+ }
}
- }
- request.Query.BoolQuery = &boolQuery
+ request.Query.BoolQuery = &boolQuery
+ }
// Add sorting if specified
if q.Sort != nil && len(*q.Sort) > 0 {
@@ -661,6 +678,10 @@ func (handler *ElasticORM) SearchWithResultItemMapper(resultArray interface{}, i
// Perform the search
searchResponse, err = handler.Client.Search(indexName, &request)
+ if err == nil && shouldRetrySearchWithoutCollapse(searchResponse, q.CollapseField) {
+ request.Collapse = nil
+ searchResponse, err = handler.Client.Search(indexName, &request)
+ }
}
// Handle search errors
diff --git a/modules/elastic/schema.go b/modules/elastic/schema.go
index a437af49f..80119c1c1 100755
--- a/modules/elastic/schema.go
+++ b/modules/elastic/schema.go
@@ -35,6 +35,7 @@ import (
"sync"
"unicode"
+ "infini.sh/framework/core/elastic"
"infini.sh/framework/core/global"
"github.com/buger/jsonparser"
@@ -124,6 +125,62 @@ func parseAnnotation(mapping []util.Annotation) string {
return json
}
+func ensureDefaultStringDynamicTemplates(mappingData map[string]interface{}) {
+ if mappingData == nil {
+ return
+ }
+ if _, ok := mappingData["dynamic_templates"]; ok {
+ return
+ }
+ mappingData["dynamic_templates"] = []interface{}{
+ util.MapStr{
+ "strings": util.MapStr{
+ "match_mapping_type": "string",
+ "mapping": util.MapStr{
+ "type": "keyword",
+ "ignore_above": 256,
+ },
+ },
+ },
+ }
+}
+
+func containsKeyDeep(value interface{}, targetKey string) bool {
+ switch v := value.(type) {
+ case map[string]interface{}:
+ for key, nested := range v {
+ if key == targetKey {
+ return true
+ }
+ if containsKeyDeep(nested, targetKey) {
+ return true
+ }
+ }
+ case []interface{}:
+ for _, nested := range v {
+ if containsKeyDeep(nested, targetKey) {
+ return true
+ }
+ }
+ }
+ return false
+}
+
+func shouldRefreshExistingTemplate(client elastic.API, templateName string, mappingData map[string]interface{}) bool {
+ if mappingData == nil {
+ return false
+ }
+ if _, ok := mappingData["dynamic_templates"]; !ok {
+ return false
+ }
+ template, err := client.GetTemplate(templateName)
+ if err != nil {
+ log.Warnf("failed to inspect existing template [%s]: %v", templateName, err)
+ return false
+ }
+ return !containsKeyDeep(template, "dynamic_templates")
+}
+
func initIndexName(t interface{}, indexName string) string {
pkg, ojbType := util.GetTypeAndPackageName(t, true)
key := fmt.Sprintf("%s-%s", pkg, ojbType)
@@ -181,6 +238,7 @@ func (handler *ElasticORM) RegisterSchemaWithName(t interface{}, indexName strin
}
return err
}
+ ensureDefaultStringDynamicTemplates(mappingData)
template, err := handler.Client.BuildTemplate(indexName+"*", nil, mappingData)
if err != nil {
if handler.Config.PanicOnInitSchemaError {
@@ -228,6 +286,44 @@ func (handler *ElasticORM) RegisterSchemaWithName(t interface{}, indexName strin
//init index
_ = handler.tryCreateInitIndex(t, indexName)
+ } else if handler.Config.BuildTemplateForObject {
+ jsonFormat := `{ %s }`
+ mapping := getIndexMapping(t)
+ js := parseAnnotation(mapping)
+ json := fmt.Sprintf(jsonFormat, quoteJson(js))
+
+ var mappingData map[string]interface{}
+ err = util.FromJSONBytes([]byte(json), &mappingData)
+ if err != nil {
+ if handler.Config.PanicOnInitSchemaError {
+ panic(err)
+ }
+ return err
+ }
+ ensureDefaultStringDynamicTemplates(mappingData)
+ if shouldRefreshExistingTemplate(handler.Client, indexTemplate, mappingData) {
+ template, err := handler.Client.BuildTemplate(indexName+"*", nil, mappingData)
+ if err != nil {
+ if handler.Config.PanicOnInitSchemaError {
+ panic(err)
+ }
+ return err
+ }
+ data, err := handler.Client.PutTemplate(indexTemplate, template)
+ if err != nil {
+ if handler.Config.PanicOnInitSchemaError {
+ panic(err)
+ }
+ return err
+ }
+ x, _, _, _ := jsonparser.Get(data, "error")
+ if x != nil {
+ log.Errorf("error on update template: %v, %v", indexName, string(x))
+ if handler.Config.PanicOnInitSchemaError {
+ panic(string(data))
+ }
+ }
+ }
}
return err
}
diff --git a/modules/elastic/schema_test.go b/modules/elastic/schema_test.go
index 518da84a1..ae2c39330 100644
--- a/modules/elastic/schema_test.go
+++ b/modules/elastic/schema_test.go
@@ -91,3 +91,19 @@ func TestQuoteWithUnderscore(t *testing.T) {
json := quoteJson(js)
assert.Equal(t, json, `{ "properties":{ "id": { "type": "keyword" },"created": { "type": "date" },"updated": { "type": "date" },"_system": { "type": "object" },"name": { "type": "keyword" } } }`)
}
+
+func TestEnsureDefaultStringDynamicTemplates(t *testing.T) {
+ mapping := map[string]interface{}{
+ "properties": map[string]interface{}{
+ "timestamp": map[string]interface{}{
+ "type": "date",
+ },
+ },
+ }
+
+ ensureDefaultStringDynamicTemplates(mapping)
+
+ templates, ok := mapping["dynamic_templates"].([]interface{})
+ assert.Equal(t, ok, true)
+ assert.Equal(t, len(templates), 1)
+}
diff --git a/modules/elastic/store.go b/modules/elastic/store.go
index ae49f3c31..9a43c6f99 100755
--- a/modules/elastic/store.go
+++ b/modules/elastic/store.go
@@ -38,6 +38,7 @@ import (
"infini.sh/framework/core/util"
"infini.sh/framework/modules/elastic/common"
"net/http"
+ "time"
)
type ElasticStore struct {
@@ -122,12 +123,16 @@ func (store *ElasticStore) GetValue(bucket string, key []byte) ([]byte, error) {
}
func (store *ElasticStore) AddValueCompress(bucket string, key []byte, value []byte) error {
+ return store.AddValueCompressWithTTL(bucket, key, value, 0)
+}
+
+func (store *ElasticStore) AddValueCompressWithTTL(bucket string, key []byte, value []byte, ttl time.Duration) error {
value, err := lz4.Encode(nil, value)
if err != nil {
log.Error("Failed to encode:", bucket, ",", key, ",", err)
return err
}
- return store.AddValue(bucket, key, value)
+ return store.AddValueWithTTL(bucket, key, value, ttl)
}
func getKey(bucket, key string) string {
@@ -135,6 +140,11 @@ func getKey(bucket, key string) string {
}
func (store *ElasticStore) AddValue(bucket string, key []byte, value []byte) error {
+ return store.AddValueWithTTL(bucket, key, value, 0)
+}
+
+func (store *ElasticStore) AddValueWithTTL(bucket string, key []byte, value []byte, ttl time.Duration) error {
+ _ = ttl
file := Blob{}
file.Content = base64.URLEncoding.EncodeToString(value)
_, err := store.Client.Index(store.Config.IndexName, "_doc", getKey(bucket, string(key)), file, "")
diff --git a/modules/metrics/elastic/elasticsearch.go b/modules/metrics/elastic/elasticsearch.go
index 54d9c8357..b5f27b12a 100644
--- a/modules/metrics/elastic/elasticsearch.go
+++ b/modules/metrics/elastic/elasticsearch.go
@@ -28,6 +28,7 @@ import (
"errors"
"fmt"
log "github.com/cihub/seelog"
+ "hash/fnv"
"infini.sh/framework/core/config"
"infini.sh/framework/core/elastic"
"infini.sh/framework/core/event"
@@ -127,6 +128,71 @@ func validateMonitorConfig(monitorConfig *elastic.TaskConfig) {
}
}
+func (m *ElasticsearchMetric) shouldCollectMetrics(v *elastic.ElasticsearchMetadata) bool {
+ if v == nil || v.Config == nil {
+ return false
+ }
+ if !v.Config.Monitored || !v.Config.Enabled {
+ return false
+ }
+ if m.IsAgentMode && v.Config.MetricCollectionMode == elastic.ModeAgentless {
+ log.Debugf("cluster [%v] is in agentless mode, skip agent-side metric collection", v.Config.Name)
+ return false
+ }
+ return true
+}
+
+func (m *ElasticsearchMetric) shouldCollectClusterLevelMetrics(v *elastic.ElasticsearchMetadata) bool {
+ return m.shouldCollectMetrics(v)
+}
+
+func (m *ElasticsearchMetric) shouldCollectNodeAndIndexMetrics(v *elastic.ElasticsearchMetadata) bool {
+ if !m.shouldCollectMetrics(v) {
+ return false
+ }
+ if !m.IsAgentMode && v.Config.MetricCollectionMode == elastic.ModeAgent {
+ log.Debugf("cluster [%v] is in agent mode, skip console-side node/index metric collection", v.Config.Name)
+ return false
+ }
+ return true
+}
+
+func getMetricTaskInitialDelay(clusterID, taskKind, interval string) string {
+ period := util.GetDurationOrDefault(interval, 10*time.Second)
+ if period <= 0 {
+ return ""
+ }
+
+ hasher := fnv.New64a()
+ _, _ = hasher.Write([]byte(clusterID))
+ _, _ = hasher.Write([]byte(":"))
+ _, _ = hasher.Write([]byte(taskKind))
+
+ offset := time.Duration(hasher.Sum64() % uint64(period))
+ if offset <= 0 {
+ return ""
+ }
+ return offset.String()
+}
+
+func getMetricTaskTimeout(interval string) time.Duration {
+ return util.GetDurationOrDefault(interval, 10*time.Second)
+}
+
+func wrapMetricCollectError(clusterName, metricName, endpoint, interval string, err error) error {
+ if errors.Is(err, context.DeadlineExceeded) {
+ return fmt.Errorf("[%s] collect %s from target cluster endpoint [%s] timed out after %s: %w", clusterName, metricName, endpoint, interval, err)
+ }
+ return fmt.Errorf("[%s] collect %s from target cluster endpoint [%s] failed: %w", clusterName, metricName, endpoint, err)
+}
+
+func wrapMetricPersistError(clusterName, metricName string, err error) error {
+ if errors.Is(err, context.DeadlineExceeded) {
+ return fmt.Errorf("[%s] persist %s to system metrics store timed out after target cluster collection succeeded: %w", clusterName, metricName, err)
+ }
+ return fmt.Errorf("[%s] persist %s to system metrics store failed after target cluster collection succeeded: %w", clusterName, metricName, err)
+}
+
func (m *ElasticsearchMetric) Collect() error {
if !m.Enabled {
return nil
@@ -192,30 +258,40 @@ func (m *ElasticsearchMetric) InitialCollectTask(k string, v *elastic.Elasticsea
m.RemoveTask(taskID)
}
}
- if !v.Config.Monitored || !v.Config.Enabled {
- log.Debugf("cluster [%v] NOT (enabled[%v] or monitored[%v] or not available[%v]), skip collect", v.Config.Name, v.Config.Enabled, v.Config.Monitored, v.IsAvailable())
+ if !m.shouldCollectMetrics(v) {
+ available := false
+ if v.Config.Enabled && v.Config.Monitored {
+ available = v.IsAvailable()
+ }
+ log.Debugf("cluster [%v] NOT eligible for metrics collection (enabled[%v], monitored[%v], mode[%v], available[%v]), skip collect", v.Config.Name, v.Config.Enabled, v.Config.Monitored, v.Config.MetricCollectionMode, available)
return true
}
if global.Env().IsDebug {
log.Debugf("run monitoring task for elasticsearch: %v - %v", k, v.Config.Name)
}
- var err error
monitorConfigs := getMonitorConfigs(v)
- if m.ClusterHealth && monitorConfigs.ClusterHealth.Enabled {
+ clusterLevelEnabled := m.shouldCollectClusterLevelMetrics(v)
+ nodeAndIndexEnabled := m.shouldCollectNodeAndIndexMetrics(v)
+ if !clusterLevelEnabled && !nodeAndIndexEnabled {
+ log.Debugf("cluster [%v] has no eligible metric collectors (mode[%v], agent_mode[%v])", v.Config.Name, v.Config.MetricCollectionMode, m.IsAgentMode)
+ return true
+ }
+ if clusterLevelEnabled && m.ClusterHealth && monitorConfigs.ClusterHealth.Enabled {
log.Debugf("collect cluster health: %s, endpoint: %s", k, v.Config.GetAnyEndpoint())
var clusterHealthMetricTask = task.ScheduleTask{
- ID: clusterHealthTaskID,
- Description: fmt.Sprintf("monitoring cluster health metric for cluster %s", k),
- Type: "interval",
- Singleton: true,
- Interval: monitorConfigs.ClusterHealth.Interval,
+ ID: clusterHealthTaskID,
+ Description: fmt.Sprintf("monitoring cluster health metric for cluster %s", k),
+ Type: "interval",
+ Singleton: true,
+ Interval: monitorConfigs.ClusterHealth.Interval,
+ InitialDelay: getMetricTaskInitialDelay(k, "cluster_health", monitorConfigs.ClusterHealth.Interval),
Task: func(ctx context.Context) {
if !v.IsAvailable() {
log.Debugf("cluster [%v] is not available, skip collect cluster health metric", v.Config.Name)
return
}
- err = m.CollectClusterHealth(k, v)
+ err := m.CollectClusterHealth(k, v)
if err != nil {
log.Error("collect cluster health error: ", err)
}
@@ -226,20 +302,21 @@ func (m *ElasticsearchMetric) InitialCollectTask(k string, v *elastic.Elasticsea
}
//cluster stats
- if m.ClusterStats && monitorConfigs.ClusterStats.Enabled {
+ if clusterLevelEnabled && m.ClusterStats && monitorConfigs.ClusterStats.Enabled {
log.Debugf("collect cluster state: %s, endpoint: %s", k, v.Config.GetAnyEndpoint())
var clusterStatsMetricTask = task.ScheduleTask{
- ID: clusterStatsTaskID,
- Description: fmt.Sprintf("monitoring cluster stats metric for cluster %s", k),
- Type: "interval",
- Singleton: true,
- Interval: monitorConfigs.ClusterStats.Interval,
+ ID: clusterStatsTaskID,
+ Description: fmt.Sprintf("monitoring cluster stats metric for cluster %s", k),
+ Type: "interval",
+ Singleton: true,
+ Interval: monitorConfigs.ClusterStats.Interval,
+ InitialDelay: getMetricTaskInitialDelay(k, "cluster_stats", monitorConfigs.ClusterStats.Interval),
Task: func(ctx context.Context) {
if !v.IsAvailable() {
log.Debugf("cluster [%v] is not available, skip collect cluster stats metric", v.Config.Name)
return
}
- err = m.CollectClusterState(k, v)
+ err := m.CollectClusterState(k, v)
if err != nil {
log.Error("collect cluster state error: ", err)
}
@@ -250,13 +327,14 @@ func (m *ElasticsearchMetric) InitialCollectTask(k string, v *elastic.Elasticsea
}
//nodes stats
- if m.NodeStats && monitorConfigs.NodeStats.Enabled {
+ if nodeAndIndexEnabled && m.NodeStats && monitorConfigs.NodeStats.Enabled {
var nodeStatsMetricTask = task.ScheduleTask{
- ID: nodeStatsTaskID,
- Description: fmt.Sprintf("monitoring node stats metric for cluster %s", k),
- Type: "interval",
- Interval: monitorConfigs.NodeStats.Interval,
- Singleton: true,
+ ID: nodeStatsTaskID,
+ Description: fmt.Sprintf("monitoring node stats metric for cluster %s", k),
+ Type: "interval",
+ Interval: monitorConfigs.NodeStats.Interval,
+ InitialDelay: getMetricTaskInitialDelay(k, "node_stats", monitorConfigs.NodeStats.Interval),
+ Singleton: true,
Task: func(ctx context.Context) {
if !v.IsAvailable() {
log.Debugf("cluster [%v] is not available, skip collect node stats metric", v.Config.Name)
@@ -267,7 +345,7 @@ func (m *ElasticsearchMetric) InitialCollectTask(k string, v *elastic.Elasticsea
)
client := elastic.GetClient(k)
- shards, err = client.CatShards()
+ shards, err := client.CatShards()
if err != nil {
log.Debug(v.Config.Name, " get shards info error: ", err)
}
@@ -312,7 +390,9 @@ func (m *ElasticsearchMetric) InitialCollectTask(k string, v *elastic.Elasticsea
if _, ok := shardInfos[nodeID]; ok {
shardInfos[nodeID]["indices_count"] = len(indexInfos[nodeID])
}
- m.SaveNodeStats(v, nodeID, nodeStats, shardInfos[nodeID])
+ if err := m.SaveNodeStats(v, nodeID, nodeStats, shardInfos[nodeID]); err != nil {
+ log.Error("collect node stats error: ", err)
+ }
}
}
} else {
@@ -326,13 +406,14 @@ func (m *ElasticsearchMetric) InitialCollectTask(k string, v *elastic.Elasticsea
}
//indices stats
- if (m.AllIndexStats || m.IndexStats) && monitorConfigs.IndexStats.Enabled {
+ if nodeAndIndexEnabled && (m.AllIndexStats || m.IndexStats) && monitorConfigs.IndexStats.Enabled {
var indexStatsMetricTask = task.ScheduleTask{
- ID: indexStatsTaskID,
- Description: fmt.Sprintf("monitoring index stats metric for cluster %s", k),
- Type: "interval",
- Interval: monitorConfigs.IndexStats.Interval,
- Singleton: true,
+ ID: indexStatsTaskID,
+ Description: fmt.Sprintf("monitoring index stats metric for cluster %s", k),
+ Type: "interval",
+ Interval: monitorConfigs.IndexStats.Interval,
+ InitialDelay: getMetricTaskInitialDelay(k, "index_stats", monitorConfigs.IndexStats.Interval),
+ Singleton: true,
Task: func(ctx context.Context) {
if !v.IsAvailable() {
log.Debugf("cluster [%v] is not available, skip collect index stats metric", v.Config.Name)
@@ -343,7 +424,7 @@ func (m *ElasticsearchMetric) InitialCollectTask(k string, v *elastic.Elasticsea
)
client := elastic.GetClient(k)
- shards, err = client.CatShards()
+ shards, err := client.CatShards()
if err != nil {
log.Debug(v.Config.Name, " get shards info error: ", err)
//return true
@@ -362,10 +443,11 @@ func (m *ElasticsearchMetric) InitialCollectTask(k string, v *elastic.Elasticsea
shardInfos := map[string][]elastic.CatShardResponse{}
if v.IsAvailable() {
- indexInfos, err = client.GetIndices("")
+ fetchedIndexInfos, err := client.GetIndices("")
if err != nil {
log.Error(v.Config.Name, " get indices info error: ", err)
}
+ indexInfos = fetchedIndexInfos
for _, item := range shards {
if _, ok := shardInfos[item.Index]; !ok {
@@ -379,7 +461,9 @@ func (m *ElasticsearchMetric) InitialCollectTask(k string, v *elastic.Elasticsea
}
if m.AllIndexStats {
- m.SaveIndexStats(v, "_all", "_all", indexStats.All.Primaries, indexStats.All.Total, nil, nil)
+ if err := m.SaveIndexStats(v, "_all", "_all", indexStats.All.Primaries, indexStats.All.Total, nil, nil); err != nil {
+ log.Error("collect index stats error: ", err)
+ }
}
if m.IndexStats {
@@ -392,7 +476,9 @@ func (m *ElasticsearchMetric) InitialCollectTask(k string, v *elastic.Elasticsea
if shardInfos != nil {
shardInfo = shardInfos[x]
}
- m.SaveIndexStats(v, y.Uuid, x, y.Primaries, y.Total, &indexInfo, shardInfo)
+ if err := m.SaveIndexStats(v, y.Uuid, x, y.Primaries, y.Total, &indexInfo, shardInfo); err != nil {
+ log.Error("collect index stats error: ", err)
+ }
}
}
}
@@ -446,7 +532,11 @@ func (m *ElasticsearchMetric) SaveNodeStats(v *elastic.ElasticsearchMetadata, no
},
}
- return m.onSaveEvent(&item)
+ if err := m.onSaveEvent(&item); err != nil {
+ return wrapMetricPersistError(v.Config.Name, fmt.Sprintf("node_stats[%s]", nodeID), err)
+ }
+
+ return nil
}
func (m *ElasticsearchMetric) SaveIndexStats(v *elastic.ElasticsearchMetadata, indexID, indexName string, primary, total elastic.IndexLevelStats, info *elastic.IndexInfo, shardInfo []elastic.CatShardResponse) error {
@@ -484,7 +574,11 @@ func (m *ElasticsearchMetric) SaveIndexStats(v *elastic.ElasticsearchMetadata, i
},
}
- return m.onSaveEvent(&item)
+ if err := m.onSaveEvent(&item); err != nil {
+ return wrapMetricPersistError(v.Config.Name, fmt.Sprintf("index_stats[%s]", indexName), err)
+ }
+
+ return nil
}
func (m *ElasticsearchMetric) CollectClusterHealth(k string, v *elastic.ElasticsearchMetadata) error {
@@ -495,7 +589,7 @@ func (m *ElasticsearchMetric) CollectClusterHealth(k string, v *elastic.Elastics
//add context to control timeout for metric collecting,
//since next metric collecting round will be triggered after this one
monitorCfg := getMonitorConfigs(v)
- du, _ := time.ParseDuration(monitorCfg.ClusterHealth.Interval)
+ du := getMetricTaskTimeout(monitorCfg.ClusterHealth.Interval)
ctx, cancel := context.WithTimeout(context.Background(), du)
defer cancel()
var (
@@ -504,13 +598,7 @@ func (m *ElasticsearchMetric) CollectClusterHealth(k string, v *elastic.Elastics
)
health, err = client.ClusterHealthSpecEndpoint(ctx, v.Config.GetAnyEndpoint(), "indices")
if err != nil {
- if errors.Is(err, context.DeadlineExceeded) {
- // Explicitly handle context deadline exceeded
- return fmt.Errorf("[%s] get cluster health context deadline exceeded after %s: %w", v.Config.Name, monitorCfg.ClusterHealth.Interval, err)
- } else {
- // Handle other errors
- return fmt.Errorf("[%s] get cluster health error: %w", v.Config.Name, err)
- }
+ return wrapMetricCollectError(v.Config.Name, "cluster_health", v.Config.GetAnyEndpoint(), monitorCfg.ClusterHealth.Interval, err)
}
indicesHealth := health.Indices
@@ -534,7 +622,7 @@ func (m *ElasticsearchMetric) CollectClusterHealth(k string, v *elastic.Elastics
err = m.onSaveEvent(&item)
if err != nil {
- return fmt.Errorf("[%s] save cluster health error: %w", v.Config.Name, err)
+ return wrapMetricPersistError(v.Config.Name, "cluster_health", err)
}
for indexName, healthInfo := range indicesHealth {
item = event.Event{
@@ -556,7 +644,7 @@ func (m *ElasticsearchMetric) CollectClusterHealth(k string, v *elastic.Elastics
}
err = m.onSaveEvent(&item)
if err != nil {
- return fmt.Errorf("[%s] save index health error: %w", v.Config.Name, err)
+ return wrapMetricPersistError(v.Config.Name, fmt.Sprintf("index_health[%s]", indexName), err)
}
}
return nil
@@ -565,6 +653,7 @@ func (m *ElasticsearchMetric) CollectClusterHealth(k string, v *elastic.Elastics
func (m *ElasticsearchMetric) CollectClusterState(k string, v *elastic.ElasticsearchMetadata) error {
log.Trace("collecting custer state metrics for :", k)
+ startAt := time.Now()
client := elastic.GetClient(k)
@@ -572,7 +661,7 @@ func (m *ElasticsearchMetric) CollectClusterState(k string, v *elastic.Elasticse
//add context to control timeout for metric collecting,
//since next metric collecting round will be triggered after this one
monitorCfg := getMonitorConfigs(v)
- du, _ := time.ParseDuration(monitorCfg.ClusterHealth.Interval)
+ du := getMetricTaskTimeout(monitorCfg.ClusterStats.Interval)
ctx, cancel := context.WithTimeout(context.Background(), du)
defer cancel()
var err error
@@ -582,13 +671,29 @@ func (m *ElasticsearchMetric) CollectClusterState(k string, v *elastic.Elasticse
stats, err = client.GetClusterStats(ctx, "")
}
if err != nil {
- if errors.Is(err, context.DeadlineExceeded) {
- // Explicitly handle context deadline exceeded
- return fmt.Errorf("[%s] get cluster stats context deadline exceeded after %s: %w", v.Config.Name, monitorCfg.ClusterHealth.Interval, err)
- } else {
- // Handle other errors
- return fmt.Errorf("[%s] get cluster stats error: %w", v.Config.Name, err)
+ return wrapMetricCollectError(v.Config.Name, "cluster_stats", v.Config.GetAnyEndpoint(), monitorCfg.ClusterStats.Interval, err)
+ }
+ elapsed := time.Since(startAt)
+ if elapsed > du*8/10 {
+ responseSize := uint64(0)
+ if stats != nil && stats.RawResult != nil {
+ responseSize = stats.RawResult.Size
+ }
+ indexFieldCount := 0
+ nodeFieldCount := 0
+ if stats != nil {
+ indexFieldCount = len(stats.Indices)
+ nodeFieldCount = len(stats.Nodes)
}
+ log.Warnf(
+ "collect cluster_stats for cluster [%s] is near timeout, elapsed: %v, timeout: %s, response_size: %d bytes, index_fields: %d, node_fields: %d",
+ v.Config.Name,
+ elapsed,
+ monitorCfg.ClusterStats.Interval,
+ responseSize,
+ indexFieldCount,
+ nodeFieldCount,
+ )
}
item := event.Event{
@@ -609,7 +714,11 @@ func (m *ElasticsearchMetric) CollectClusterState(k string, v *elastic.Elasticse
},
}
- return m.onSaveEvent(&item)
+ if err := m.onSaveEvent(&item); err != nil {
+ return wrapMetricPersistError(v.Config.Name, "cluster_stats", err)
+ }
+
+ return nil
}
func (m *ElasticsearchMetric) CollectNodeStats() {
diff --git a/modules/metrics/elastic/elasticsearch_test.go b/modules/metrics/elastic/elasticsearch_test.go
new file mode 100644
index 000000000..d3168a9a4
--- /dev/null
+++ b/modules/metrics/elastic/elasticsearch_test.go
@@ -0,0 +1,108 @@
+package elastic
+
+import (
+ "testing"
+ "time"
+
+ coreelastic "infini.sh/framework/core/elastic"
+)
+
+func TestShouldCollectMetricsAllowsConsoleCollectorInAgentModeForClusterLevelMetrics(t *testing.T) {
+ collector := &ElasticsearchMetric{}
+ meta := &coreelastic.ElasticsearchMetadata{
+ Config: &coreelastic.ElasticsearchConfig{
+ Name: "agent-cluster",
+ Enabled: true,
+ Monitored: true,
+ MetricCollectionMode: coreelastic.ModeAgent,
+ },
+ }
+
+ if !collector.shouldCollectMetrics(meta) {
+ t.Fatal("expected console-side collector to keep cluster-level metrics in agent mode")
+ }
+ if collector.shouldCollectNodeAndIndexMetrics(meta) {
+ t.Fatal("expected console-side collector to skip node/index metrics in agent mode")
+ }
+}
+
+func TestShouldCollectMetricsAllowsConsoleCollectorInAgentlessMode(t *testing.T) {
+ collector := &ElasticsearchMetric{}
+ meta := &coreelastic.ElasticsearchMetadata{
+ Config: &coreelastic.ElasticsearchConfig{
+ Name: "agentless-cluster",
+ Enabled: true,
+ Monitored: true,
+ MetricCollectionMode: coreelastic.ModeAgentless,
+ },
+ }
+
+ if !collector.shouldCollectMetrics(meta) {
+ t.Fatal("expected console-side collector to run for agentless clusters")
+ }
+ if !collector.shouldCollectNodeAndIndexMetrics(meta) {
+ t.Fatal("expected console-side collector to run node/index metrics for agentless clusters")
+ }
+}
+
+func TestShouldCollectMetricsSkipsAgentCollectorInAgentlessMode(t *testing.T) {
+ collector := &ElasticsearchMetric{IsAgentMode: true}
+ meta := &coreelastic.ElasticsearchMetadata{
+ Config: &coreelastic.ElasticsearchConfig{
+ Name: "agentless-cluster",
+ Enabled: true,
+ Monitored: true,
+ MetricCollectionMode: coreelastic.ModeAgentless,
+ },
+ }
+
+ if collector.shouldCollectMetrics(meta) {
+ t.Fatal("expected agent-side collector to skip agentless clusters")
+ }
+}
+
+func TestShouldCollectMetricsAllowsAgentCollectorInAgentMode(t *testing.T) {
+ collector := &ElasticsearchMetric{IsAgentMode: true}
+ meta := &coreelastic.ElasticsearchMetadata{
+ Config: &coreelastic.ElasticsearchConfig{
+ Name: "agent-cluster",
+ Enabled: true,
+ Monitored: true,
+ MetricCollectionMode: coreelastic.ModeAgent,
+ },
+ }
+
+ if !collector.shouldCollectMetrics(meta) {
+ t.Fatal("expected agent-side collector to run for agent mode clusters")
+ }
+}
+
+func TestGetMetricTaskInitialDelayStableAndBounded(t *testing.T) {
+ first := getMetricTaskInitialDelay("cluster-a", "cluster_health", "10s")
+ second := getMetricTaskInitialDelay("cluster-a", "cluster_health", "10s")
+ if first != second {
+ t.Fatalf("expected stable delay, got %s and %s", first, second)
+ }
+
+ delay, err := time.ParseDuration(first)
+ if err != nil {
+ t.Fatalf("expected parseable delay, got %q: %v", first, err)
+ }
+ if delay < 0 || delay >= 10*time.Second {
+ t.Fatalf("expected delay to be within interval, got %s", delay)
+ }
+}
+
+func TestGetMetricTaskInitialDelayVariesByTaskKind(t *testing.T) {
+ healthDelay := getMetricTaskInitialDelay("cluster-a", "cluster_health", "10s")
+ statsDelay := getMetricTaskInitialDelay("cluster-a", "cluster_stats", "10s")
+ if healthDelay == statsDelay {
+ t.Fatalf("expected different metric kinds to spread across interval, both got %s", healthDelay)
+ }
+}
+
+func TestGetMetricTaskTimeoutFallsBackToDefault(t *testing.T) {
+ if got := getMetricTaskTimeout("invalid"); got != 10*time.Second {
+ t.Fatalf("expected default timeout, got %s", got)
+ }
+}
diff --git a/modules/metrics/host/overall/overall.go b/modules/metrics/host/overall/overall.go
index 81fed49c0..074980988 100644
--- a/modules/metrics/host/overall/overall.go
+++ b/modules/metrics/host/overall/overall.go
@@ -48,6 +48,7 @@ type Metric struct {
IntervalSeconds float64 `config:"interval_seconds"`
YellowThreshold float64 `config:"yellow_threshold"`
RedThreshold float64 `config:"red_threshold"`
+ event.EventSink
mu sync.Mutex
@@ -99,6 +100,11 @@ type deviceUtilization struct {
}
func New(cfg *config.Config) (*Metric, error) {
+ return NewWithSink(cfg, event.DefaultEventSink)
+}
+
+// NewWithSink creates an overall metric collector with a custom sink.
+func NewWithSink(cfg *config.Config, sink event.EventSink) (*Metric, error) {
me := &Metric{
Enabled: true,
IntervalSeconds: 10,
@@ -108,6 +114,7 @@ func New(cfg *config.Config) (*Metric, error) {
prevNetIO: make(map[string]*netIOSnapshot),
netBandwidth: make(map[string]float64),
}
+ me.EventSink = sink
err := cfg.Unpack(&me)
if err != nil {
@@ -235,7 +242,7 @@ func (m *Metric) Collect() error {
fields["status"] = status
fields["bottleneck"] = bottleneck
- return event.Save(&event.Event{
+ return m.Save(&event.Event{
Metadata: event.EventMetadata{
Category: "host",
Name: "overall",
diff --git a/modules/metrics/metrics.go b/modules/metrics/metrics.go
index 78ec1e05d..ebfbd2d6f 100755
--- a/modules/metrics/metrics.go
+++ b/modules/metrics/metrics.go
@@ -117,7 +117,8 @@ func (module *MetricsModule) Setup() {
// check other conditions
hasChanged = meta.IsAvailable() != oldMeta.IsAvailable() ||
meta.Config.Enabled != oldMeta.Config.Enabled ||
- meta.Config.Monitored != oldMeta.Config.Monitored
+ meta.Config.Monitored != oldMeta.Config.Monitored ||
+ meta.Config.MetricCollectionMode != oldMeta.Config.MetricCollectionMode
}
if !hasChanged {
return
@@ -173,7 +174,7 @@ func (module *MetricsModule) CollectAgentMetric() {
Type: "interval",
Interval: "10s",
Task: func(ctx context.Context) {
- log.Debug("collecting instance metrics")
+ log.Trace("collecting instance metrics")
agentM.Collect()
},
}
@@ -204,7 +205,7 @@ func (module *MetricsModule) CollectHostMetric() {
Type: "interval",
Interval: "10s",
Task: func(ctx context.Context) {
- log.Debug("collecting network metrics")
+ log.Trace("collecting network metrics")
netM.Collect()
},
}
diff --git a/modules/pipeline/model.go b/modules/pipeline/model.go
index 10c4b52f5..0fce2368b 100644
--- a/modules/pipeline/model.go
+++ b/modules/pipeline/model.go
@@ -30,12 +30,21 @@ import (
"infini.sh/framework/core/util"
)
-type PipelineTaskStatus struct {
- State pipeline.RunningState `json:"state"`
- CreateTime time.Time `json:"create_time"`
- StartTime *time.Time `json:"start_time"`
- EndTime *time.Time `json:"end_time"`
- Context util.MapStr `json:"context"`
- Config *pipeline.PipelineConfigV2 `json:"config"`
- Processors []map[string]interface{} `json:"processor"`
+type PipelineStatus struct {
+ State pipeline.RunningState `json:"state"`
+ LastRunState pipeline.RunningState `json:"last_run_state,omitempty"`
+ CreateTime time.Time `json:"create_time"`
+ StartTime *time.Time `json:"start_time"`
+ EndTime *time.Time `json:"end_time"`
+ Context util.MapStr `json:"context"`
+ Result *PipelineResult `json:"result,omitempty"`
+ Config *pipeline.PipelineConfigV2 `json:"config"`
+ Processors []map[string]interface{} `json:"processor"`
+}
+
+type PipelineTaskStatus = PipelineStatus
+
+type PipelineResult struct {
+ Success bool `json:"success"`
+ Error string `json:"error,omitempty"`
}
diff --git a/modules/pipeline/pipeline_test.go b/modules/pipeline/pipeline_test.go
new file mode 100644
index 000000000..a34b123ea
--- /dev/null
+++ b/modules/pipeline/pipeline_test.go
@@ -0,0 +1,50 @@
+package pipeline
+
+import (
+ "testing"
+ "time"
+
+ corepipeline "infini.sh/framework/core/pipeline"
+)
+
+func TestDeleteTaskWaitsForLoopRelease(t *testing.T) {
+ module := &PipeModule{}
+ ctx := corepipeline.AcquireContext(corepipeline.PipelineConfigV2{})
+
+ module.contexts.Store("task-1", ctx)
+ module.configs.Store("task-1", corepipeline.PipelineConfigV2{Name: "task-1"})
+ module.pipelines.Store("task-1", struct{}{})
+
+ released := make(chan struct{})
+ go func() {
+ for !ctx.IsCanceled() {
+ time.Sleep(time.Millisecond)
+ }
+ time.Sleep(50 * time.Millisecond)
+ ctx.SetLoopReleased()
+ close(released)
+ }()
+
+ start := time.Now()
+ module.deleteTask("task-1")
+ elapsed := time.Since(start)
+
+ select {
+ case <-released:
+ default:
+ t.Fatal("expected deleteTask to wait for loop release")
+ }
+
+ if elapsed < 50*time.Millisecond {
+ t.Fatalf("expected deleteTask to wait for loop release, returned after %v", elapsed)
+ }
+ if _, ok := module.contexts.Load("task-1"); ok {
+ t.Fatal("expected context to be deleted")
+ }
+ if _, ok := module.configs.Load("task-1"); ok {
+ t.Fatal("expected config to be deleted")
+ }
+ if _, ok := module.pipelines.Load("task-1"); ok {
+ t.Fatal("expected pipeline to be deleted")
+ }
+}
diff --git a/modules/pipeline/proto.go b/modules/pipeline/proto.go
index 922bb27cc..7bc29ca39 100644
--- a/modules/pipeline/proto.go
+++ b/modules/pipeline/proto.go
@@ -25,7 +25,9 @@ package pipeline
import "infini.sh/framework/core/pipeline"
-type GetPipelineTasksResponse map[string]*PipelineTaskStatus
+type GetPipelinesResponse map[string]*PipelineStatus
+
+type GetPipelineTasksResponse = GetPipelinesResponse
type CreatePipelineRequest struct {
pipeline.PipelineConfigV2
diff --git a/modules/pipeline/tasks.go b/modules/pipeline/tasks.go
index eb7bc6662..b516de9e1 100644
--- a/modules/pipeline/tasks.go
+++ b/modules/pipeline/tasks.go
@@ -79,12 +79,19 @@ func (module *PipeModule) getPipelineTaskStatus(id string, config string, proces
if !ok {
return nil
}
- ret := &PipelineTaskStatus{
- State: c1.GetRunningState(),
- CreateTime: c1.GetCreateTime(),
- StartTime: c1.GetStartTime(),
- EndTime: c1.GetEndTime(),
- Context: c1.CloneData(),
+ ret := &PipelineStatus{
+ State: c1.GetRunningState(),
+ LastRunState: c1.GetResultState(),
+ CreateTime: c1.GetCreateTime(),
+ StartTime: c1.GetStartTime(),
+ EndTime: c1.GetEndTime(),
+ Context: c1.CloneData(),
+ }
+ if ret.LastRunState == pipeline.FINISHED || ret.LastRunState == pipeline.FAILED {
+ ret.Result = &PipelineResult{
+ Success: c1.GetResultError() == "",
+ Error: c1.GetResultError(),
+ }
}
if config != "false" {
v1, ok := module.configs.Load(id)
diff --git a/modules/queue/disk_queue/cleanup.go b/modules/queue/disk_queue/cleanup.go
index 1fed12942..3456d881f 100644
--- a/modules/queue/disk_queue/cleanup.go
+++ b/modules/queue/disk_queue/cleanup.go
@@ -92,7 +92,7 @@ func (module *DiskQueue) deleteUnusedFiles(queueID string, fileNum int64) {
fileStartToDelete := fileNum - module.cfg.Retention.MaxNumOfLocalFiles
if fileStartToDelete <= 0 || consumers <= 0 || eSegmentNum < 0 {
- log.Debugf("queue: %v, no consumers or consumer/s3 already ahead of this file, %v, %v, %v", queueID, fileStartToDelete, consumers, eSegmentNum)
+ log.Tracef("queue: %v, no consumers or consumer/s3 already ahead of this file, %v, %v, %v", queueID, fileStartToDelete, consumers, eSegmentNum)
return
}
diff --git a/modules/queue/disk_queue/compress.go b/modules/queue/disk_queue/compress.go
index b03fdccfb..451e66e19 100644
--- a/modules/queue/disk_queue/compress.go
+++ b/modules/queue/disk_queue/compress.go
@@ -34,6 +34,7 @@ import (
"infini.sh/framework/core/util"
"infini.sh/framework/core/util/zstd"
"os"
+ "strings"
"sync"
)
@@ -104,7 +105,7 @@ func (module *DiskQueue) compressFiles(queueID string, fileNum int64) {
//skip compress file
if fileStartToCompress <= 0 || (module.cfg.SkipZeroConsumers && consumers <= 0) || fileStartToCompress <= lastCompressedFileNum {
- log.Debugf("skip compress %v", queueID)
+ log.Tracef("skip compress %v", queueID)
return
}
@@ -124,6 +125,10 @@ func (module *DiskQueue) compressFiles(queueID string, fileNum int64) {
//compress
err := zstd.CompressFile(file, toFile)
if err != nil {
+ if strings.Contains(err.Error(), "temp file for target file was exits, skip:") {
+ log.Debug(err)
+ continue
+ }
log.Error(err)
continue
}
diff --git a/modules/queue/disk_queue/consumer.go b/modules/queue/disk_queue/consumer.go
index c18bbd83e..cf0d14c75 100644
--- a/modules/queue/disk_queue/consumer.go
+++ b/modules/queue/disk_queue/consumer.go
@@ -68,6 +68,29 @@ type Consumer struct {
fileLoadCompleted bool
}
+func (d *Consumer) parkOnEmptyTail(fileName string) error {
+ if d.readFile != nil {
+ if err := d.readFile.Close(); err != nil && !util.ContainStr(err.Error(), "already") {
+ return err
+ }
+ }
+ d.readFile = nil
+ d.reader = nil
+ d.fileName = fileName
+ d.lastFileSize = 0
+ d.maxBytesPerFileRead = 0
+ d.fileLoadCompleted = false
+ return nil
+}
+
+func (d *Consumer) waitingForTailFile() bool {
+ return d.diskQueue != nil &&
+ d.readFile == nil &&
+ d.reader == nil &&
+ d.segment == d.diskQueue.writeSegmentNum &&
+ d.readPos == 0
+}
+
func (c *Consumer) getFileSize() int64 {
var err error
readFile, err := os.OpenFile(c.fileName, os.O_RDONLY, 0600)
@@ -144,6 +167,22 @@ READ_MSG:
// check reader
if d.reader == nil {
+ if d.waitingForTailFile() {
+ if d.diskQueue.writePos > 0 || util.FileExists(d.fileName) {
+ err = d.ResetOffset(d.segment, d.readPos)
+ if err != nil {
+ if strings.Contains(err.Error(), "not found") {
+ return messages, false, nil
+ }
+ return messages, false, err
+ }
+ goto READ_MSG
+ }
+ if len(messages) == 0 && d.cCfg.EOFRetryDelayInMs > 0 {
+ time.Sleep(time.Duration(d.cCfg.EOFRetryDelayInMs) * time.Millisecond)
+ }
+ return messages, false, nil
+ }
return messages, false, errors.New("reader is nil")
}
//read message size
@@ -206,7 +245,7 @@ READ_MSG:
oldPart := d.segment
Notify(d.queue, ReadComplete, d.segment)
ctx.UpdateNextOffset(d.segment, d.readPos) //update next offset
- log.Debugf("EOF, but current read segment_id [%v] is less than current write segment_id [%v], increase ++", oldPart, d.diskQueue.writeSegmentNum)
+ log.Tracef("EOF, but current read segment_id [%v] is less than current write segment_id [%v], increase ++", oldPart, d.diskQueue.writeSegmentNum)
err = d.ResetOffset(d.segment+1, 0) //locate next segment
if err != nil {
if strings.Contains(err.Error(), "not found") {
@@ -236,7 +275,7 @@ READ_MSG:
}
return messages, false, err
}
- log.Debugf("queue:%v, offset:%v,%v, msgSize:%v", d.queue, d.segment, d.readPos, msgSize)
+ log.Tracef("queue:%v, offset:%v,%v, msgSize:%v", d.queue, d.segment, d.readPos, msgSize)
if int32(msgSize) < d.mCfg.MinMsgSize || int32(msgSize) > d.mCfg.MaxMsgSize {
//current have changes, reload file with new position
newFileSize := d.getFileSize()
@@ -274,8 +313,19 @@ READ_MSG:
//can't read ahead before current write file
if nextSegment >= d.diskQueue.writeSegmentNum {
log.Debugf("need to skip to next file, but next file not exists, current write segment:%v, current read segment:%v", d.diskQueue.writeSegmentNum, d.segment)
- d.diskQueue.skipToNextRWFile(false)
+ err = d.diskQueue.skipToNextRWFile(false)
+ if err != nil {
+ return messages, false, err
+ }
d.diskQueue.needSync = true
+ err = d.ResetOffset(d.diskQueue.writeSegmentNum, 0)
+ if err != nil {
+ if strings.Contains(err.Error(), "not found") {
+ return messages, false, nil
+ }
+ return messages, false, err
+ }
+ ctx.UpdateNextOffset(d.segment, d.readPos)
} else {
//let's continue move to next file
nextSegment++
@@ -332,9 +382,9 @@ READ_MSG:
//still working on the same file
if d.diskQueue.writeSegmentNum == d.segment {
time.Sleep(100 * time.Millisecond) // Prevent catching up too quickly.
- log.Debugf("invalid message size detected. this might be due to a dirty read as the file was being written while open. reloading segment: %d", d.segment)
+ log.Tracef("invalid message size detected. this might be due to a dirty read as the file was being written while open. reloading segment: %d", d.segment)
} else {
- log.Debugf("invalid message size detected. this might be due to a partial file load. reloading segment: %d", d.segment)
+ log.Tracef("invalid message size detected. this might be due to a partial file load. reloading segment: %d", d.segment)
}
d.readPos = previousPos
@@ -509,6 +559,9 @@ func (d *Consumer) ResetOffset(segment, readPos int64) error {
if !exists {
//double check, but next file exists
if !util.FileExists(fileName) {
+ if segment == d.diskQueue.writeSegmentNum && readPos == 0 && d.diskQueue.writePos == 0 {
+ return d.parkOnEmptyTail(fileName)
+ }
if d.mCfg.AutoSkipCorruptFile {
nextSegment := d.segment + 1
if nextSegment > d.diskQueue.writeSegmentNum {
@@ -518,7 +571,7 @@ func (d *Consumer) ResetOffset(segment, readPos int64) error {
d.qCfg.Name, d.queue, d.cCfg.Key(), d.segment, d.readPos, fileName)
RETRY_NEXT_FILE:
// there are segments in the middle
- if nextSegment < d.diskQueue.writeSegmentNum {
+ if nextSegment <= d.diskQueue.writeSegmentNum {
fileName, exists, next_file_exists = SmartGetFileName(d.mCfg, d.queue, nextSegment)
if exists || util.FileExists(fileName) {
log.Debugf("retry skip to next file: %v, exists", fileName)
@@ -532,6 +585,12 @@ func (d *Consumer) ResetOffset(segment, readPos int64) error {
goto RETRY_NEXT_FILE
}
} else {
+ if d.diskQueue.writePos == 0 {
+ d.segment = d.diskQueue.writeSegmentNum
+ d.readPos = 0
+ d.diskQueue.UpdateSegmentConsumerInReading(d.ID, d.segment)
+ return d.parkOnEmptyTail(GetFileName(d.queue, d.segment))
+ }
return errors.New(fileName + " not found, next segment greater than current write segment")
}
} else {
diff --git a/modules/queue/disk_queue/diskqueue.go b/modules/queue/disk_queue/diskqueue.go
index dd9db9236..8e67506a8 100644
--- a/modules/queue/disk_queue/diskqueue.go
+++ b/modules/queue/disk_queue/diskqueue.go
@@ -69,6 +69,8 @@ import (
"infini.sh/framework/core/util/zstd"
)
+const bytesPerMiB = 1024 * 1024
+
// providing a filesystem backed FIFO queue
type DiskBasedQueue struct {
sync.RWMutex
@@ -118,6 +120,8 @@ type DiskBasedQueue struct {
// NewDiskQueue instantiates a new instance of DiskBasedQueue, retrieving metadata
// from the filesystem and starting the read ahead goroutine
func NewDiskQueueByConfig(name, dataPath string, cfg *DiskQueueConfig) *DiskBasedQueue {
+ normalizeDiskQueueConfig(cfg)
+
d := DiskBasedQueue{
name: name,
dataPath: dataPath,
@@ -139,6 +143,9 @@ func NewDiskQueueByConfig(name, dataPath string, cfg *DiskQueueConfig) *DiskBase
if err != nil && !os.IsNotExist(err) {
log.Errorf("diskqueue(%s) failed to retrieveMetaData - %s", d.name, err)
}
+ if repairErr := d.repairTailMetadata(); repairErr != nil {
+ log.Errorf("diskqueue(%s) failed to repair tail metadata - %s", d.name, repairErr)
+ }
// Always advance to a new segment on process restart to prevent data loss
// or corruption from overwriting existing segment files.
@@ -197,7 +204,8 @@ func (d *DiskBasedQueue) ReadChan() <-chan []byte {
// Put writes a []byte to the queue
func (d *DiskBasedQueue) Put(data []byte) WriteResponse {
- ctx, cancel := context.WithTimeout(context.Background(), time.Duration(d.cfg.WriteTimeoutInMS)*time.Millisecond)
+ writeTimeout := d.getWriteTimeout(len(data))
+ ctx, cancel := context.WithTimeout(context.Background(), writeTimeout)
defer cancel()
size := int64(len(data))
@@ -252,7 +260,7 @@ func (d *DiskBasedQueue) Put(data []byte) WriteResponse {
switch res.Error {
case context.DeadlineExceeded:
// Handle timeout error specifically
- res.Error = fmt.Errorf("operation timed out: %w", res.Error)
+ res.Error = fmt.Errorf("operation timed out after %s waiting for disk queue writer availability: %w", writeTimeout, res.Error)
case context.Canceled:
// Handle cancellation error specifically
res.Error = fmt.Errorf("operation was canceled: %w", res.Error)
@@ -264,13 +272,31 @@ func (d *DiskBasedQueue) Put(data []byte) WriteResponse {
}
}
+func (d *DiskBasedQueue) getWriteTimeout(payloadSize int) time.Duration {
+ timeoutInMS := defaultWriteTimeoutInMS
+ if d != nil && d.cfg != nil && d.cfg.WriteTimeoutInMS > 0 {
+ timeoutInMS = d.cfg.WriteTimeoutInMS
+ }
+
+ if payloadSize > 0 {
+ payloadMiB := int64((payloadSize + bytesPerMiB - 1) / bytesPerMiB)
+ timeoutInMS += payloadMiB * adaptiveWriteTimeoutPerPayloadMiBInMS
+ }
+
+ if d != nil && len(d.writeChan) > 0 {
+ timeoutInMS += int64(len(d.writeChan)) * adaptiveWriteTimeoutPerQueuedWriteInMS
+ }
+
+ if timeoutInMS > maxAdaptiveWriteTimeoutInMS {
+ timeoutInMS = maxAdaptiveWriteTimeoutInMS
+ }
+
+ return time.Duration(timeoutInMS) * time.Millisecond
+}
+
// Close cleans up the queue and persists metadata
func (d *DiskBasedQueue) Close() error {
- err := d.exit(false)
- if err != nil {
- return err
- }
- return d.sync()
+ return d.exit(false)
}
// Destroy cleans up all data for the specified queue
@@ -297,6 +323,13 @@ func (d *DiskBasedQueue) Delete() error {
return d.exit(true)
}
+func (d *DiskBasedQueue) ensureDataPathExists() error {
+ if d == nil || d.dataPath == "" {
+ return nil
+ }
+ return os.MkdirAll(d.dataPath, 0o755)
+}
+
func (d *DiskBasedQueue) exit(deleted bool) error {
d.Lock()
@@ -330,6 +363,11 @@ func (d *DiskBasedQueue) exit(deleted bool) error {
// ensure that ioLoop has exited
<-d.exitSyncChan
+ var syncErr error
+ if !deleted {
+ syncErr = d.sync()
+ }
+
close(d.depthChan)
if d.readFile != nil {
@@ -342,7 +380,7 @@ func (d *DiskBasedQueue) exit(deleted bool) error {
d.writeFile = nil
}
- return nil
+ return syncErr
}
// Empty destructively clears out any pending data in the queue
@@ -559,6 +597,11 @@ func (d *DiskBasedQueue) writeOne(data []byte) WriteResponse {
var res WriteResponse
if d.writeFile == nil {
+ err = d.ensureDataPathExists()
+ if err != nil {
+ res.Error = err
+ return res
+ }
curFileName := d.GetFileName(d.writeSegmentNum)
d.writeFile, err = os.OpenFile(curFileName, os.O_RDWR|os.O_CREATE, 0600)
if err != nil {
@@ -705,6 +748,158 @@ func (d *DiskBasedQueue) retrieveMetaData() error {
return nil
}
+type segmentScanResult struct {
+ validEnd int64
+ totalMessages int64
+ messagesBeforeReadPos int64
+ messagesBeforeWritePos int64
+ readBoundary int64
+}
+
+func scanSegmentFileTail(file *os.File, cfg *DiskQueueConfig, readPos, writePos int64) (segmentScanResult, error) {
+ result := segmentScanResult{}
+ if file == nil || cfg == nil {
+ return result, nil
+ }
+
+ if _, err := file.Seek(0, 0); err != nil {
+ return result, err
+ }
+
+ reader := bufio.NewReader(file)
+ var offset int64
+
+ for {
+ var msgSize int32
+ if err := binary.Read(reader, binary.BigEndian, &msgSize); err != nil {
+ if err == io.EOF || err == io.ErrUnexpectedEOF {
+ return result, nil
+ }
+ return result, err
+ }
+
+ if msgSize < cfg.MinMsgSize || msgSize > cfg.MaxMsgSize {
+ return result, nil
+ }
+
+ payloadSize := int64(msgSize)
+ if _, err := io.CopyN(io.Discard, reader, payloadSize); err != nil {
+ if err == io.EOF || err == io.ErrUnexpectedEOF {
+ return result, nil
+ }
+ return result, err
+ }
+
+ offset += 4 + payloadSize
+ result.validEnd = offset
+ result.totalMessages++
+ if offset <= readPos {
+ result.messagesBeforeReadPos++
+ result.readBoundary = offset
+ }
+ if offset <= writePos {
+ result.messagesBeforeWritePos++
+ }
+ }
+}
+
+func (d *DiskBasedQueue) repairTailMetadata() error {
+ if d == nil || d.cfg == nil {
+ return nil
+ }
+ if d.writeSegmentNum == 0 && d.writePos == 0 {
+ return nil
+ }
+
+ fileName := d.GetFileName(d.writeSegmentNum)
+ if !util.FileExists(fileName) {
+ return nil
+ }
+
+ file, err := os.OpenFile(fileName, os.O_RDWR, 0600)
+ if err != nil {
+ return err
+ }
+ defer file.Close()
+
+ stat, err := file.Stat()
+ if err != nil {
+ return err
+ }
+
+ readPos := int64(0)
+ if d.readSegmentFileNum == d.writeSegmentNum {
+ readPos = d.readPos
+ }
+
+ scan, err := scanSegmentFileTail(file, d.cfg, readPos, d.writePos)
+ if err != nil {
+ return err
+ }
+
+ newWritePos := scan.validEnd
+ newReadPos := readPos
+ if d.readSegmentFileNum == d.writeSegmentNum {
+ if newReadPos > newWritePos {
+ newReadPos = newWritePos
+ }
+ if scan.readBoundary < newReadPos {
+ newReadPos = scan.readBoundary
+ }
+ }
+
+ oldUnreadInTail := scan.messagesBeforeWritePos
+ newUnreadInTail := scan.totalMessages
+ if d.readSegmentFileNum == d.writeSegmentNum {
+ oldUnreadInTail -= scan.messagesBeforeReadPos
+ newUnreadInTail -= scan.messagesBeforeReadPos
+ }
+
+ changed := false
+ if stat.Size() != scan.validEnd {
+ if err := file.Truncate(scan.validEnd); err != nil {
+ return err
+ }
+ if err := file.Sync(); err != nil {
+ return err
+ }
+ log.Warnf("diskqueue(%s) truncated tail segment %s from %d to %d bytes during startup recovery",
+ d.name, fileName, stat.Size(), scan.validEnd)
+ changed = true
+ }
+
+ if d.writePos != newWritePos {
+ d.writePos = newWritePos
+ changed = true
+ }
+
+ if d.readSegmentFileNum == d.writeSegmentNum && d.readPos != newReadPos {
+ d.readPos = newReadPos
+ d.nextReadPos = newReadPos
+ changed = true
+ }
+
+ if delta := newUnreadInTail - oldUnreadInTail; delta != 0 {
+ d.depth += delta
+ if d.depth < 0 {
+ d.depth = 0
+ }
+ changed = true
+ }
+
+ if d.nextReadFileNum == d.writeSegmentNum && d.nextReadPos > d.writePos {
+ d.nextReadPos = d.writePos
+ changed = true
+ }
+
+ if !changed {
+ return nil
+ }
+
+ d.needSync = true
+ return d.persistMetaData()
+}
+
// persistMetaData atomically writes state to the filesystem
func (d *DiskBasedQueue) persistMetaData() error {
d.metaLock.Lock()
@@ -729,6 +924,11 @@ func (d *DiskBasedQueue) persistMetaData() error {
var f *os.File
var err error
+ err = d.ensureDataPathExists()
+ if err != nil {
+ return err
+ }
+
fileName := d.metaDataFileName()
tmpFileName := fmt.Sprintf("%s.%d.tmp", fileName, rand.Int())
diff --git a/modules/queue/disk_queue/diskqueue_test.go b/modules/queue/disk_queue/diskqueue_test.go
new file mode 100644
index 000000000..d041b8ab5
--- /dev/null
+++ b/modules/queue/disk_queue/diskqueue_test.go
@@ -0,0 +1,449 @@
+package queue
+
+import (
+ "encoding/binary"
+ "os"
+ "path/filepath"
+ "sync"
+ "testing"
+ "time"
+
+ . "infini.sh/framework/core/env"
+ "infini.sh/framework/core/global"
+ corequeue "infini.sh/framework/core/queue"
+)
+
+func TestGetWriteTimeoutIncludesPayloadAndBacklog(t *testing.T) {
+ dq := &DiskBasedQueue{
+ cfg: &DiskQueueConfig{WriteTimeoutInMS: defaultWriteTimeoutInMS},
+ writeChan: make(chan []byte, defaultWriteChanBuffer),
+ }
+
+ dq.writeChan <- []byte("a")
+ dq.writeChan <- []byte("b")
+
+ timeout := dq.getWriteTimeout(3 * bytesPerMiB)
+
+ expected := time.Duration(defaultWriteTimeoutInMS+3*adaptiveWriteTimeoutPerPayloadMiBInMS+2*adaptiveWriteTimeoutPerQueuedWriteInMS) * time.Millisecond
+ if timeout != expected {
+ t.Fatalf("unexpected write timeout: got %s want %s", timeout, expected)
+ }
+}
+
+func TestGetWriteTimeoutCapsAtMaximum(t *testing.T) {
+ dq := &DiskBasedQueue{
+ cfg: &DiskQueueConfig{WriteTimeoutInMS: defaultWriteTimeoutInMS},
+ writeChan: make(chan []byte, defaultWriteChanBuffer),
+ }
+
+ for i := 0; i < cap(dq.writeChan); i++ {
+ dq.writeChan <- []byte("x")
+ }
+
+ timeout := dq.getWriteTimeout(64 * bytesPerMiB)
+ expected := time.Duration(maxAdaptiveWriteTimeoutInMS) * time.Millisecond
+ if timeout != expected {
+ t.Fatalf("unexpected capped timeout: got %s want %s", timeout, expected)
+ }
+}
+
+func TestClosePersistsUnsyncedWritesBeforeClosingFiles(t *testing.T) {
+ env1 := EmptyEnv()
+ env1.SystemConfig.PathConfig.Data = t.TempDir()
+ global.RegisterEnv(env1)
+
+ queueName := "close-persists-unsynced"
+ cfg := &DiskQueueConfig{
+ MinMsgSize: 1,
+ MaxMsgSize: 1024,
+ MaxBytesPerFile: 1024 * 1024,
+ SyncEveryRecords: 1 << 20,
+ SyncTimeoutInMS: 1 << 20,
+ ReadChanBuffer: 0,
+ WriteChanBuffer: 1,
+ }
+ normalizeDiskQueueConfig(cfg)
+
+ dataPath := GetDataPath(queueName)
+ if err := os.MkdirAll(dataPath, 0o755); err != nil {
+ t.Fatalf("failed to create queue data dir: %v", err)
+ }
+
+ dq := &DiskBasedQueue{
+ name: queueName,
+ dataPath: dataPath,
+ cfg: cfg,
+ readChan: make(chan []byte, cfg.ReadChanBuffer),
+ depthChan: make(chan int64),
+ writeChan: make(chan []byte, cfg.WriteChanBuffer),
+ writeResponseChan: make(chan WriteResponse),
+ emptyChan: make(chan int),
+ emptyResponseChan: make(chan error),
+ exitChan: make(chan int),
+ exitSyncChan: make(chan int, 1),
+ consumersInReading: sync.Map{},
+ }
+ go dq.ioLoop()
+
+ res := dq.Put([]byte("hello"))
+ if res.Error != nil {
+ t.Fatalf("failed to put queue message: %v", res.Error)
+ }
+ if dq.writeFile == nil {
+ t.Fatalf("expected queue write file to remain open before close")
+ }
+
+ if err := dq.Close(); err != nil {
+ t.Fatalf("failed to close queue: %v", err)
+ }
+
+ reopened := &DiskBasedQueue{
+ name: queueName,
+ dataPath: dataPath,
+ cfg: cfg,
+ readChan: make(chan []byte, cfg.ReadChanBuffer),
+ depthChan: make(chan int64),
+ writeChan: make(chan []byte, cfg.WriteChanBuffer),
+ writeResponseChan: make(chan WriteResponse),
+ emptyChan: make(chan int),
+ emptyResponseChan: make(chan error),
+ exitChan: make(chan int),
+ exitSyncChan: make(chan int, 1),
+ consumersInReading: sync.Map{},
+ }
+ if err := reopened.retrieveMetaData(); err != nil {
+ t.Fatalf("failed to reload queue metadata: %v", err)
+ }
+ t.Cleanup(func() {
+ _ = os.RemoveAll(dataPath)
+ })
+
+ if depth := reopened.depth; depth != 1 {
+ t.Fatalf("expected reopened queue depth 1, got %d", depth)
+ }
+
+ message, err := reopened.readOne()
+ if err != nil {
+ t.Fatalf("failed to read reopened queue message: %v", err)
+ }
+ if string(message) != "hello" {
+ t.Fatalf("expected reopened queue payload %q, got %q", "hello", string(message))
+ }
+}
+
+func TestRepairTailMetadataTruncatesIncompleteTailOnStartup(t *testing.T) {
+ env1 := EmptyEnv()
+ env1.SystemConfig.PathConfig.Data = t.TempDir()
+ global.RegisterEnv(env1)
+
+ queueName := "repair-tail-startup"
+ cfg := &DiskQueueConfig{
+ MinMsgSize: 1,
+ MaxMsgSize: 1024,
+ MaxBytesPerFile: 1024 * 1024,
+ }
+ normalizeDiskQueueConfig(cfg)
+
+ fileName := GetFileName(queueName, 0)
+ if err := os.MkdirAll(filepath.Dir(fileName), 0o755); err != nil {
+ t.Fatalf("failed to create queue dir: %v", err)
+ }
+
+ payload := []byte("hello")
+ file, err := os.Create(fileName)
+ if err != nil {
+ t.Fatalf("failed to create tail file: %v", err)
+ }
+ if err := binary.Write(file, binary.BigEndian, int32(len(payload))); err != nil {
+ t.Fatalf("failed to write payload size: %v", err)
+ }
+ if _, err := file.Write(payload); err != nil {
+ t.Fatalf("failed to write payload: %v", err)
+ }
+ if _, err := file.Write([]byte{0x7f, 0xff}); err != nil {
+ t.Fatalf("failed to append corrupt tail: %v", err)
+ }
+ if err := file.Close(); err != nil {
+ t.Fatalf("failed to close tail file: %v", err)
+ }
+
+ dq := &DiskBasedQueue{
+ name: queueName,
+ dataPath: GetDataPath(queueName),
+ cfg: cfg,
+ readSegmentFileNum: 0,
+ writeSegmentNum: 0,
+ readPos: 0,
+ nextReadPos: 0,
+ writePos: int64(4 + len(payload) + 2),
+ depth: 1,
+ }
+
+ if err := dq.repairTailMetadata(); err != nil {
+ t.Fatalf("failed to repair tail metadata: %v", err)
+ }
+
+ expectedWritePos := int64(4 + len(payload))
+ if dq.writePos != expectedWritePos {
+ t.Fatalf("expected write pos %d after repair, got %d", expectedWritePos, dq.writePos)
+ }
+ if dq.depth != 1 {
+ t.Fatalf("expected queue depth to remain 1 after repair, got %d", dq.depth)
+ }
+
+ stat, err := os.Stat(fileName)
+ if err != nil {
+ t.Fatalf("failed to stat repaired tail file: %v", err)
+ }
+ if stat.Size() != expectedWritePos {
+ t.Fatalf("expected repaired tail file size %d, got %d", expectedWritePos, stat.Size())
+ }
+
+ message, err := dq.readOne()
+ if err != nil {
+ t.Fatalf("failed to read message after repair: %v", err)
+ }
+ if string(message) != "hello" {
+ t.Fatalf("expected repaired payload %q, got %q", "hello", string(message))
+ }
+}
+
+func TestQueueRecreatesDataPathAfterDirectoryDeletion(t *testing.T) {
+ env1 := EmptyEnv()
+ env1.SystemConfig.PathConfig.Data = t.TempDir()
+ global.RegisterEnv(env1)
+
+ queueName := "recreate-data-path"
+ cfg := &DiskQueueConfig{
+ MinMsgSize: 1,
+ MaxMsgSize: 1024,
+ MaxBytesPerFile: 1024 * 1024,
+ }
+ normalizeDiskQueueConfig(cfg)
+
+ dataPath := GetDataPath(queueName)
+ if err := os.MkdirAll(dataPath, 0o755); err != nil {
+ t.Fatalf("failed to create queue data dir: %v", err)
+ }
+
+ dq := &DiskBasedQueue{
+ name: queueName,
+ dataPath: dataPath,
+ cfg: cfg,
+ writePos: 0,
+ writeFile: nil,
+ }
+
+ if err := os.RemoveAll(dataPath); err != nil {
+ t.Fatalf("failed to delete queue data dir: %v", err)
+ }
+
+ if err := dq.sync(); err != nil {
+ t.Fatalf("expected sync to recreate deleted queue dir, got %v", err)
+ }
+ if _, err := os.Stat(dataPath); err != nil {
+ t.Fatalf("expected queue data dir to be recreated, got %v", err)
+ }
+
+ res := dq.writeOne([]byte("hello"))
+ if res.Error != nil {
+ t.Fatalf("expected write to succeed after dir recreation, got %v", res.Error)
+ }
+ if _, err := os.Stat(dq.metaDataFileName()); err != nil {
+ t.Fatalf("expected metadata file to be recreated, got %v", err)
+ }
+ if _, err := os.Stat(dq.GetFileName(0)); err != nil {
+ t.Fatalf("expected segment file to be recreated, got %v", err)
+ }
+
+ if dq.writeFile != nil {
+ _ = dq.writeFile.Close()
+ }
+}
+
+func TestCloseSucceedsAfterQueueDirectoryDeletion(t *testing.T) {
+ env1 := EmptyEnv()
+ env1.SystemConfig.PathConfig.Data = t.TempDir()
+ global.RegisterEnv(env1)
+
+ queueName := "close-after-delete"
+ cfg := &DiskQueueConfig{
+ MinMsgSize: 1,
+ MaxMsgSize: 1024,
+ MaxBytesPerFile: 1024 * 1024,
+ SyncEveryRecords: 1 << 20,
+ SyncTimeoutInMS: 1 << 20,
+ ReadChanBuffer: 0,
+ WriteChanBuffer: 1,
+ }
+ normalizeDiskQueueConfig(cfg)
+
+ dataPath := GetDataPath(queueName)
+ if err := os.MkdirAll(dataPath, 0o755); err != nil {
+ t.Fatalf("failed to create queue data dir: %v", err)
+ }
+
+ dq := &DiskBasedQueue{
+ name: queueName,
+ dataPath: dataPath,
+ cfg: cfg,
+ readChan: make(chan []byte, cfg.ReadChanBuffer),
+ depthChan: make(chan int64),
+ writeChan: make(chan []byte, cfg.WriteChanBuffer),
+ writeResponseChan: make(chan WriteResponse),
+ emptyChan: make(chan int),
+ emptyResponseChan: make(chan error),
+ exitChan: make(chan int),
+ exitSyncChan: make(chan int, 1),
+ consumersInReading: sync.Map{},
+ }
+ go dq.ioLoop()
+
+ res := dq.Put([]byte("hello"))
+ if res.Error != nil {
+ t.Fatalf("failed to put queue message: %v", res.Error)
+ }
+
+ if err := os.RemoveAll(dataPath); err != nil {
+ t.Fatalf("failed to delete queue data dir: %v", err)
+ }
+
+ if err := dq.Close(); err != nil {
+ t.Fatalf("expected close to succeed after queue dir deletion, got %v", err)
+ }
+}
+
+func TestResetOffsetSkipsMissingSegmentsUpToCurrentWriteSegment(t *testing.T) {
+ env1 := EmptyEnv()
+ env1.SystemConfig.PathConfig.Data = t.TempDir()
+ global.RegisterEnv(env1)
+
+ queueName := "reset-offset-skip"
+ data := []byte("ok")
+ fileName := GetFileName(queueName, 2)
+ if err := os.MkdirAll(filepath.Dir(fileName), 0o755); err != nil {
+ t.Fatalf("failed to create queue dir: %v", err)
+ }
+ file, err := os.Create(fileName)
+ if err != nil {
+ t.Fatalf("failed to create segment file: %v", err)
+ }
+ if err := binary.Write(file, binary.BigEndian, int32(len(data))); err != nil {
+ t.Fatalf("failed to write message size: %v", err)
+ }
+ if _, err := file.Write(data); err != nil {
+ t.Fatalf("failed to write message body: %v", err)
+ }
+ if err := file.Close(); err != nil {
+ t.Fatalf("failed to close segment file: %v", err)
+ }
+
+ dq := &DiskBasedQueue{
+ name: queueName,
+ cfg: &DiskQueueConfig{AutoSkipCorruptFile: true, MinMsgSize: 1, MaxMsgSize: 1024},
+ writeSegmentNum: 2,
+ writePos: int64(4 + len(data)),
+ }
+ consumer := &Consumer{
+ ID: "consumer-reset",
+ diskQueue: dq,
+ mCfg: dq.cfg,
+ qCfg: &corequeue.QueueConfig{Name: queueName},
+ cCfg: &corequeue.ConsumerConfig{},
+ queue: queueName,
+ }
+
+ if err := consumer.ResetOffset(1, 0); err != nil {
+ t.Fatalf("expected reset offset to skip to current write segment, got %v", err)
+ }
+ if consumer.segment != 2 {
+ t.Fatalf("expected consumer to move to segment 2, got %d", consumer.segment)
+ }
+ if consumer.reader == nil {
+ t.Fatalf("expected consumer reader to be initialized for segment 2")
+ }
+}
+
+func TestFetchMessagesRecoversToEmptyTailWithoutRescanningCorruptFile(t *testing.T) {
+ env1 := EmptyEnv()
+ env1.SystemConfig.PathConfig.Data = t.TempDir()
+ global.RegisterEnv(env1)
+
+ queueName := "fetch-empty-tail"
+ corruptFile := GetFileName(queueName, 1)
+ if err := os.MkdirAll(filepath.Dir(corruptFile), 0o755); err != nil {
+ t.Fatalf("failed to create queue dir: %v", err)
+ }
+ if err := os.WriteFile(corruptFile, []byte{0x7f, 0xff, 0xff, 0xff}, 0o644); err != nil {
+ t.Fatalf("failed to write corrupt segment: %v", err)
+ }
+
+ dq := &DiskBasedQueue{
+ name: queueName,
+ cfg: &DiskQueueConfig{AutoSkipCorruptFile: true, MinMsgSize: 1, MaxMsgSize: 1024},
+ writeSegmentNum: 3,
+ writePos: 0,
+ }
+ consumer := &Consumer{
+ ID: "consumer-fetch",
+ diskQueue: dq,
+ mCfg: dq.cfg,
+ qCfg: &corequeue.QueueConfig{Name: queueName},
+ cCfg: &corequeue.ConsumerConfig{},
+ queue: queueName,
+ }
+
+ if err := consumer.ResetOffset(1, 0); err != nil {
+ t.Fatalf("failed to initialize consumer: %v", err)
+ }
+
+ ctx := &corequeue.Context{}
+ messages, timeout, err := consumer.FetchMessages(ctx, 1)
+ if err != nil {
+ t.Fatalf("expected corruption recovery without error, got %v", err)
+ }
+ if timeout {
+ t.Fatalf("did not expect timeout during corruption recovery")
+ }
+ if len(messages) != 0 {
+ t.Fatalf("expected no messages during recovery, got %d", len(messages))
+ }
+ if consumer.segment != dq.writeSegmentNum {
+ t.Fatalf("expected consumer to park on new tail segment %d, got %d", dq.writeSegmentNum, consumer.segment)
+ }
+ if ctx.NextOffset.Segment != dq.writeSegmentNum || ctx.NextOffset.Position != 0 {
+ t.Fatalf("expected next offset to advance to new tail, got %v", ctx.NextOffset)
+ }
+
+ payload := []byte("hello")
+ tailFile := GetFileName(queueName, dq.writeSegmentNum)
+ file, err := os.Create(tailFile)
+ if err != nil {
+ t.Fatalf("failed to create new tail segment: %v", err)
+ }
+ if err := binary.Write(file, binary.BigEndian, int32(len(payload))); err != nil {
+ t.Fatalf("failed to write tail message size: %v", err)
+ }
+ if _, err := file.Write(payload); err != nil {
+ t.Fatalf("failed to write tail message body: %v", err)
+ }
+ if err := file.Close(); err != nil {
+ t.Fatalf("failed to close tail segment: %v", err)
+ }
+ dq.writePos = int64(4 + len(payload))
+
+ ctx = &corequeue.Context{}
+ messages, timeout, err = consumer.FetchMessages(ctx, 1)
+ if err != nil {
+ t.Fatalf("expected consumer to resume reading on new tail, got %v", err)
+ }
+ if timeout {
+ t.Fatalf("did not expect timeout when new tail data exists")
+ }
+ if len(messages) != 1 {
+ t.Fatalf("expected exactly one message, got %d", len(messages))
+ }
+ if string(messages[0].Data) != "hello" {
+ t.Fatalf("expected payload %q, got %q", "hello", string(messages[0].Data))
+ }
+}
diff --git a/modules/queue/disk_queue/module.go b/modules/queue/disk_queue/module.go
index d7b1d8ee9..2383a68b4 100644
--- a/modules/queue/disk_queue/module.go
+++ b/modules/queue/disk_queue/module.go
@@ -124,8 +124,33 @@ type CompressConfig struct {
Level int `config:"level"`
}
+const (
+ defaultWriteTimeoutInMS int64 = 60 * 1000
+ defaultWriteChanBuffer = 16
+ minRecommendedWriteTimeoutInMS int64 = 15 * 1000
+ maxAdaptiveWriteTimeoutInMS int64 = 5 * 60 * 1000
+ adaptiveWriteTimeoutPerQueuedWriteInMS int64 = 3 * 1000
+ adaptiveWriteTimeoutPerPayloadMiBInMS int64 = 5 * 1000
+)
+
var preventRead bool
+func normalizeDiskQueueConfig(cfg *DiskQueueConfig) {
+ if cfg == nil {
+ return
+ }
+
+ if cfg.WriteTimeoutInMS <= 0 {
+ cfg.WriteTimeoutInMS = defaultWriteTimeoutInMS
+ } else if cfg.WriteTimeoutInMS < minRecommendedWriteTimeoutInMS {
+ log.Warnf("disk_queue write timeout may be too small on slow disks: %dms", cfg.WriteTimeoutInMS)
+ }
+
+ if cfg.WriteChanBuffer <= 0 {
+ cfg.WriteChanBuffer = defaultWriteChanBuffer
+ }
+}
+
func checkCapacity(cfg *DiskQueueConfig) error {
if cfg.CheckDiskCapacityRetryDelayInMs <= 0 {
@@ -233,20 +258,20 @@ func (module *DiskQueue) Setup() {
MinMsgSize: 1,
MaxMsgSize: 104857600, //100MB
MaxBytesPerFile: 100 * 1024 * 1024, //100MB
- WriteTimeoutInMS: 1000, //1s
- CheckDiskCapacityRetryDelayInMs: 10 * 000, //10s
+ WriteTimeoutInMS: defaultWriteTimeoutInMS,
+ CheckDiskCapacityRetryDelayInMs: 10 * 000, //10s
EOFRetryDelayInMs: 500,
SyncEveryRecords: 1000,
SyncTimeoutInMS: 1000,
NotifyChanBuffer: 100,
ReadChanBuffer: 0,
- WriteChanBuffer: 0,
+ WriteChanBuffer: defaultWriteChanBuffer,
WarningFreeBytes: 10 * 1024 * 1024 * 1024,
ReservedFreeBytes: 5 * 1024 * 1024 * 1024,
PrepareFilesToRead: true,
Compress: DiskCompress{
IdleThreshold: 3,
- DeleteAfterCompress: false,
+ DeleteAfterCompress: true,
NumOfFilesDecompressAhead: 3,
Message: CompressConfig{
Enabled: false,
@@ -262,6 +287,8 @@ func (module *DiskQueue) Setup() {
panic(err)
}
+ normalizeDiskQueueConfig(module.cfg)
+
if !module.cfg.Enabled {
return
}
diff --git a/modules/queue/disk_queue/module_test.go b/modules/queue/disk_queue/module_test.go
new file mode 100644
index 000000000..fc643621f
--- /dev/null
+++ b/modules/queue/disk_queue/module_test.go
@@ -0,0 +1,49 @@
+package queue
+
+import (
+ "testing"
+
+ . "infini.sh/framework/core/env"
+ "infini.sh/framework/core/global"
+)
+
+func TestNormalizeDiskQueueConfigAppliesRobustWriteDefaults(t *testing.T) {
+ cfg := &DiskQueueConfig{}
+
+ normalizeDiskQueueConfig(cfg)
+
+ if cfg.WriteTimeoutInMS != defaultWriteTimeoutInMS {
+ t.Fatalf("unexpected write timeout: %d", cfg.WriteTimeoutInMS)
+ }
+ if cfg.WriteChanBuffer != defaultWriteChanBuffer {
+ t.Fatalf("unexpected write chan buffer: %d", cfg.WriteChanBuffer)
+ }
+}
+
+func TestNormalizeDiskQueueConfigKeepsExplicitWriteSettings(t *testing.T) {
+ cfg := &DiskQueueConfig{
+ WriteTimeoutInMS: 45 * 1000,
+ WriteChanBuffer: 64,
+ }
+
+ normalizeDiskQueueConfig(cfg)
+
+ if cfg.WriteTimeoutInMS != 45*1000 {
+ t.Fatalf("write timeout should be preserved, got %d", cfg.WriteTimeoutInMS)
+ }
+ if cfg.WriteChanBuffer != 64 {
+ t.Fatalf("write chan buffer should be preserved, got %d", cfg.WriteChanBuffer)
+ }
+}
+
+func TestSetupDefaultsDeleteAfterCompress(t *testing.T) {
+ env1 := EmptyEnv()
+ global.RegisterEnv(env1)
+
+ module := DiskQueue{}
+ module.Setup()
+
+ if !module.cfg.Compress.DeleteAfterCompress {
+ t.Fatalf("delete_after_compress should default to true")
+ }
+}
diff --git a/modules/security/access_token/authentication.go b/modules/security/access_token/authentication.go
index a1c860ee8..6c0808778 100644
--- a/modules/security/access_token/authentication.go
+++ b/modules/security/access_token/authentication.go
@@ -22,11 +22,12 @@ import (
"infini.sh/framework/core/orm"
"infini.sh/framework/core/security"
"infini.sh/framework/core/util"
- "infini.sh/framework/modules/security/http_filters"
)
const ProviderName = "access_token"
+const defaultAPITokenTTL = 365 * 24 * time.Hour
+
const (
// KVAccessTokenBucket stores token_string -> AccessToken JSON. Used by
// byAPITokenHeader to authenticate inbound requests in both modes.
@@ -68,7 +69,7 @@ func init() {
security.RegisterHTTPAuthFilterProviderWithPriority("api_token", byAPITokenHeader, 30)
api.HandleUIMethod(api.POST, "/auth/access_token", RequestAccessToken, api.RequirePermission(createTokenPermission))
- api.HandleUIMethod(api.GET, "/auth/access_token/_search", SearchAccessToken, api.RequirePermission(searchTokenPermission), api.Feature(http_filters.FeatureMaskSensitiveField))
+ api.HandleUIMethod(api.GET, "/auth/access_token/_search", SearchAccessToken, api.RequirePermission(searchTokenPermission))
api.HandleUIMethod(api.DELETE, "/auth/access_token/:token_id", DeleteAccessToken, api.RequirePermission(deleteTokenPermission))
api.HandleUIMethod(api.PUT, "/auth/access_token/:token_id", UpdateAccessToken, api.RequirePermission(updateTokenPermission))
@@ -172,6 +173,7 @@ func RequestAccessToken(w http.ResponseWriter, req *http.Request, ps httprouter.
reqBody := struct {
Name string `json:"name"`
Description string `json:"description"`
+ ExpireIn *int64 `json:"expire_in,omitempty"`
Permissions []security.PermissionKey `json:"permissions,omitempty"`
}{}
err = api.DecodeJSON(req, &reqBody)
@@ -202,7 +204,10 @@ func RequestAccessToken(w http.ResponseWriter, req *http.Request, ps httprouter.
}
}
- expiredAT := time.Now().Add(365 * 24 * time.Hour).Unix()
+ expiredAT, err := normalizeAPITokenExpireAt(reqBody.ExpireIn)
+ if err != nil {
+ panic(errors.ErrorWithHTTPCode(err, 400, "invalid expire_in"))
+ }
res, err := CreateAPIToken(reqUser, reqBody.Name, reqBody.Description, "general", expiredAT, permissions)
if err != nil {
panic(err)
@@ -410,6 +415,7 @@ func UpdateAccessToken(w http.ResponseWriter, req *http.Request, ps httprouter.P
reqBody := struct {
Name string `json:"name,omitempty"`
Description string `json:"description"`
+ ExpireIn *int64 `json:"expire_in,omitempty"`
Permissions []security.PermissionKey `json:"permissions,omitempty"`
}{}
err = api.DecodeJSON(req, &reqBody)
@@ -459,6 +465,13 @@ func UpdateAccessToken(w http.ResponseWriter, req *http.Request, ps httprouter.P
if reqBody.Description != "" {
token.Description = reqBody.Description
}
+ if reqBody.ExpireIn != nil {
+ expiredAT, err := normalizeAPITokenExpireAt(reqBody.ExpireIn)
+ if err != nil {
+ panic(errors.ErrorWithHTTPCode(err, 400, "invalid expire_in"))
+ }
+ token.ExpireIn = expiredAT
+ }
if len(reqBody.Permissions) > 0 {
if isNative() {
@@ -491,6 +504,19 @@ func UpdateAccessToken(w http.ResponseWriter, req *http.Request, ps httprouter.P
api.WriteUpdatedOKJSON(w, tokenID)
}
+func normalizeAPITokenExpireAt(expireIn *int64) (int64, error) {
+ if expireIn == nil {
+ return time.Now().Add(defaultAPITokenTTL).Unix(), nil
+ }
+ if *expireIn <= 0 {
+ return 0, nil
+ }
+ if *expireIn <= time.Now().Unix() {
+ return 0, errors.Errorf("expire_in must be greater than current time")
+ }
+ return *expireIn, nil
+}
+
// GenerateApiTokenName generates a unique API token name
func GenerateApiTokenName(prefix string) string {
if prefix == "" {
diff --git a/modules/security/account/profile.go b/modules/security/account/profile.go
index f8f9f0692..0895e613f 100644
--- a/modules/security/account/profile.go
+++ b/modules/security/account/profile.go
@@ -28,10 +28,21 @@ func Profile(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
}
p := &security.UserProfile{
- Name: reqUser.Login,
+ Name: reqUser.Login,
+ Roles: reqUser.Roles,
}
p.ID = reqUser.UserID
+ if reqUser.Provider == security.DefaultNativeAuthBackend {
+ if _, account, err := security.GetUserByID(reqUser.UserID); err == nil && account != nil {
+ if account.Name != "" {
+ p.Name = account.Name
+ }
+ p.Email = account.Email
+ p.Roles = account.Roles
+ }
+ }
+
//get all permissions for user
p.Permissions = reqUser.GetPermissionKeys()
diff --git a/modules/security/account/refresh.go b/modules/security/account/refresh.go
new file mode 100644
index 000000000..5d1f51fa7
--- /dev/null
+++ b/modules/security/account/refresh.go
@@ -0,0 +1,102 @@
+/* Copyright © INFINI LTD. All rights reserved.
+ * Web: https://infinilabs.com
+ * Email: hello#infini.ltd */
+
+package account
+
+import (
+ "fmt"
+ "net/http"
+ "strings"
+
+ "infini.sh/framework/core/api"
+ httprouter "infini.sh/framework/core/api/router"
+ "infini.sh/framework/core/security"
+)
+
+func init() {
+ api.HandleUIMethod(
+ api.POST,
+ "/account/refresh",
+ api.RequireSecureTransport(Refresh, api.SecureTransportOptions{TrustForwardHeaders: true}),
+ api.RequireLogin(),
+ api.AllowOPTIONSS(),
+ api.Feature(api.FeatureCORS),
+ )
+}
+
+// Refresh reissues an access token for the current session user while reloading the
+// native account record so updated roles/profile data are reflected in new tokens.
+func Refresh(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ reqUser, err := security.GetUserFromContext(r.Context())
+ if err != nil || reqUser == nil {
+ api.WriteError(w, "invalid user", http.StatusUnauthorized)
+ return
+ }
+
+ sessionUser, err := buildRefreshedSession(reqUser)
+ if err != nil {
+ api.WriteError(w, err.Error(), http.StatusUnauthorized)
+ return
+ }
+
+ if err, token := security.AddUserToSession(w, r, sessionUser); err != nil {
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ } else {
+ security.DecorateSessionTokenResponse(token, sessionUser)
+ api.WriteOKJSON(w, token)
+ }
+}
+
+func buildRefreshedSession(reqUser *security.UserSessionInfo) (*security.UserSessionInfo, error) {
+ if reqUser == nil {
+ return nil, fmt.Errorf("user not found")
+ }
+
+ sessionUser := cloneSessionUser(reqUser)
+ if reqUser.Provider != security.DefaultNativeAuthBackend {
+ return sessionUser, nil
+ }
+
+ provider, account, err := security.GetUserByID(reqUser.UserID)
+ if err != nil {
+ return nil, err
+ }
+ if account == nil {
+ return nil, fmt.Errorf("user not found")
+ }
+
+ login := strings.TrimSpace(account.Email)
+ if login == "" {
+ login = strings.TrimSpace(reqUser.Login)
+ }
+ if provider == "" {
+ provider = security.DefaultNativeAuthBackend
+ }
+
+ sessionUser = &security.UserSessionInfo{
+ Provider: provider,
+ Login: login,
+ Roles: append([]string(nil), account.Roles...),
+ LastLogin: reqUser.LastLogin,
+ }
+ sessionUser.SetUserID(account.ID)
+ sessionUser.UserAssignedPermission = security.NewUserAssignedPermission(security.GetAllPermissionsForUser(sessionUser), nil)
+ return sessionUser, nil
+}
+
+func cloneSessionUser(reqUser *security.UserSessionInfo) *security.UserSessionInfo {
+ if reqUser == nil {
+ return nil
+ }
+
+ sessionUser := &security.UserSessionInfo{
+ Provider: reqUser.Provider,
+ Login: reqUser.Login,
+ Roles: append([]string(nil), reqUser.Roles...),
+ LastLogin: reqUser.LastLogin,
+ }
+ sessionUser.SetUserID(reqUser.UserID)
+ sessionUser.UserAssignedPermission = security.NewUserAssignedPermission(security.GetAllPermissionsForUser(sessionUser), nil)
+ return sessionUser
+}
diff --git a/modules/security/account/refresh_test.go b/modules/security/account/refresh_test.go
new file mode 100644
index 000000000..41e49aecb
--- /dev/null
+++ b/modules/security/account/refresh_test.go
@@ -0,0 +1,98 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package account
+
+import (
+ "testing"
+
+ "infini.sh/framework/core/security"
+)
+
+type refreshTestProvider struct{}
+
+func (refreshTestProvider) GetUserByID(id string) (bool, *security.UserAccount, error) {
+ if id != "refresh-native-user" {
+ return false, nil, nil
+ }
+
+ account := &security.UserAccount{
+ Name: "Refreshed Admin",
+ Email: "refreshed@example.org",
+ Roles: []string{security.RoleAdmin},
+ }
+ account.ID = id
+ return true, account, nil
+}
+
+func (refreshTestProvider) GetUserByLogin(login string) (bool, *security.UserAccount, error) {
+ return false, nil, nil
+}
+
+func (refreshTestProvider) CreateUser(name, login, password string, force bool) (*security.UserAccount, error) {
+ return nil, nil
+}
+
+// External providers can keep their current session payload when refreshing.
+func TestBuildRefreshedSessionKeepsExternalUser(t *testing.T) {
+ reqUser := &security.UserSessionInfo{
+ Provider: "sso",
+ Login: "alice@example.org",
+ Roles: []string{"viewer"},
+ }
+ reqUser.SetUserID("external-1")
+
+ sessionUser, err := buildRefreshedSession(reqUser)
+ if err != nil {
+ t.Fatalf("build refreshed session: %v", err)
+ }
+ if sessionUser.Login != reqUser.Login {
+ t.Fatalf("expected external login %q, got %q", reqUser.Login, sessionUser.Login)
+ }
+ if sessionUser.UserID != reqUser.UserID {
+ t.Fatalf("expected external user id %q, got %q", reqUser.UserID, sessionUser.UserID)
+ }
+}
+
+// Native refreshes should pull the latest account snapshot from the registered backend.
+func TestBuildRefreshedSessionReloadsNativeAccount(t *testing.T) {
+ security.RegisterAuthenticationProvider("refresh-test-provider", refreshTestProvider{})
+
+ reqUser := &security.UserSessionInfo{
+ Provider: security.DefaultNativeAuthBackend,
+ Login: "stale@example.org",
+ Roles: []string{"viewer"},
+ }
+ reqUser.SetUserID("refresh-native-user")
+
+ sessionUser, err := buildRefreshedSession(reqUser)
+ if err != nil {
+ t.Fatalf("build refreshed session: %v", err)
+ }
+ if sessionUser.Login != "refreshed@example.org" {
+ t.Fatalf("expected refreshed login, got %q", sessionUser.Login)
+ }
+ if len(sessionUser.Roles) != 1 || sessionUser.Roles[0] != security.RoleAdmin {
+ t.Fatalf("expected refreshed roles, got %#v", sessionUser.Roles)
+ }
+}
diff --git a/modules/security/http_filters/json_mask.go b/modules/security/http_filters/json_mask.go
index f98f88f6a..998b4d924 100644
--- a/modules/security/http_filters/json_mask.go
+++ b/modules/security/http_filters/json_mask.go
@@ -18,11 +18,13 @@ const FeatureRemoveSensitiveField = "feature_sensitive_fields_remove_sensitive_f
const SensitiveFields = "feature_sensitive_fields_extra_keys"
var sensitiveFields = map[string]bool{
- "password": true,
- "token": true,
- "secret": true,
- "access_token": true,
- "refresh_token": true,
+ "password": true,
+ "password_salt": true,
+ "password_verifier": true,
+ "token": true,
+ "secret": true,
+ "access_token": true,
+ "refresh_token": true,
}
type JSONMaskFilter struct{}
diff --git a/modules/security/http_filters/logging.go b/modules/security/http_filters/logging.go
index 50bfbf3db..0c5d9efc7 100644
--- a/modules/security/http_filters/logging.go
+++ b/modules/security/http_filters/logging.go
@@ -56,6 +56,10 @@ func getAccessLogHandler() *rotate.RotateWriter {
return accessLogHandler
}
+func isAccessLogEnabled() bool {
+ return global.Env().SystemConfig != nil && global.Env().SystemConfig.WebAppConfig.AccessLog
+}
+
func (f *LoggingFilter) GetPriority() int {
// Lower priority values execute first (higher precedence)
return 0
@@ -68,6 +72,10 @@ func (f *LoggingFilter) ApplyFilter(
next httprouter.Handle,
) httprouter.Handle {
return func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ if !isAccessLogEnabled() {
+ next(w, r, ps)
+ return
+ }
start := time.Now()
diff --git a/modules/security/http_filters/logging_test.go b/modules/security/http_filters/logging_test.go
new file mode 100644
index 000000000..d3d6303f7
--- /dev/null
+++ b/modules/security/http_filters/logging_test.go
@@ -0,0 +1,67 @@
+package http_filters
+
+import (
+ "net/http"
+ "net/http/httptest"
+ "os"
+ "path/filepath"
+ "testing"
+
+ httprouter "infini.sh/framework/core/api/router"
+ "infini.sh/framework/core/env"
+ "infini.sh/framework/core/global"
+)
+
+func TestLoggingFilterSkipsAccessLogWhenDisabled(t *testing.T) {
+ oldEnv := global.Env()
+ testEnv := env.EmptyEnv()
+ testEnv.SystemConfig.PathConfig.Data = t.TempDir()
+ testEnv.SystemConfig.PathConfig.Log = t.TempDir()
+ testEnv.SystemConfig.WebAppConfig.AccessLog = false
+ global.RegisterEnv(testEnv)
+ defer global.RegisterEnv(oldEnv)
+
+ accessLogHandler = nil
+ defer func() { accessLogHandler = nil }()
+
+ filter := &LoggingFilter{}
+ handler := filter.ApplyFilter(http.MethodGet, "/hello", nil, func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ w.WriteHeader(http.StatusNoContent)
+ })
+
+ req := httptest.NewRequest(http.MethodGet, "/hello", nil)
+ resp := httptest.NewRecorder()
+ handler(resp, req, nil)
+
+ accessLogPath := filepath.Join(testEnv.GetLogDir(), "access.log")
+ if _, err := os.Stat(accessLogPath); !os.IsNotExist(err) {
+ t.Fatalf("expected access log file to be absent when disabled, stat err=%v", err)
+ }
+}
+
+func TestLoggingFilterWritesAccessLogWhenEnabled(t *testing.T) {
+ oldEnv := global.Env()
+ testEnv := env.EmptyEnv()
+ testEnv.SystemConfig.PathConfig.Data = t.TempDir()
+ testEnv.SystemConfig.PathConfig.Log = t.TempDir()
+ testEnv.SystemConfig.WebAppConfig.AccessLog = true
+ global.RegisterEnv(testEnv)
+ defer global.RegisterEnv(oldEnv)
+
+ accessLogHandler = nil
+ defer func() { accessLogHandler = nil }()
+
+ filter := &LoggingFilter{}
+ handler := filter.ApplyFilter(http.MethodGet, "/hello", nil, func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ w.WriteHeader(http.StatusNoContent)
+ })
+
+ req := httptest.NewRequest(http.MethodGet, "/hello", nil)
+ resp := httptest.NewRecorder()
+ handler(resp, req, nil)
+
+ accessLogPath := filepath.Join(testEnv.GetLogDir(), "access.log")
+ if _, err := os.Stat(accessLogPath); err != nil {
+ t.Fatalf("expected access log file to exist when enabled, got %v", err)
+ }
+}
diff --git a/modules/security/http_filters/security.go b/modules/security/http_filters/security.go
new file mode 100644
index 000000000..6e372a000
--- /dev/null
+++ b/modules/security/http_filters/security.go
@@ -0,0 +1,75 @@
+/* Copyright © INFINI LTD. All rights reserved.
+ * Web: https://infinilabs.com
+ * Email: hello#infini.ltd */
+
+package http_filters
+
+import (
+ "net/http"
+
+ log "github.com/cihub/seelog"
+ "infini.sh/framework/core/api"
+ httprouter "infini.sh/framework/core/api/router"
+ replaysecurity "infini.sh/framework/core/security/replay"
+)
+
+func init() {
+ api.RegisterUIFilter(&SecurityFilter{})
+}
+
+// SecurityFilter enforces per-route HTTPS and replay-protection features declared in HandlerOptions.
+type SecurityFilter struct {
+ api.Handler
+}
+
+// GetPriority keeps the security checks ahead of permission checks but after early request shaping.
+func (f *SecurityFilter) GetPriority() int {
+ return 450
+}
+
+// ApplyFilter translates route feature flags into runtime checks for HTTPS and replay nonce usage.
+func (f *SecurityFilter) ApplyFilter(
+ method string,
+ pattern string,
+ options *api.HandlerOptions,
+ next httprouter.Handle,
+) httprouter.Handle {
+ if options == nil || (!options.Feature(api.FeatureRequireSecureTransport) && !options.Feature(api.FeatureRequireReplayProtection)) {
+ log.Debug(method, ",", pattern, ", skip security feature filters")
+ return next
+ }
+
+ return func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ if options.Feature(api.FeatureRequireSecureTransport) {
+ secureOptions := api.SecureTransportOptions{
+ TrustForwardHeaders: trustForwardHeadersFromOptions(options),
+ }
+ if !api.RequestUsesSecureTransport(r, secureOptions) {
+ f.WriteError(w, "this endpoint requires HTTPS. use https:// directly or route through a trusted HTTPS reverse proxy", http.StatusUpgradeRequired)
+ return
+ }
+ }
+
+ if options.Feature(api.FeatureRequireReplayProtection) {
+ if err := replaysecurity.ValidateAndConsumeReplayNonce(r); err != nil {
+ f.WriteError(w, err.Error(), http.StatusUnauthorized)
+ return
+ }
+ }
+
+ next(w, r, ps)
+ }
+}
+
+// trustForwardHeadersFromOptions extracts whether SecureTransportOption opted into proxy headers.
+func trustForwardHeadersFromOptions(options *api.HandlerOptions) bool {
+ if options == nil || options.Labels == nil {
+ return false
+ }
+ trustValue, ok := options.Labels[api.LabelTrustForwardHeaders]
+ if !ok {
+ return false
+ }
+ trustForwardHeaders, ok := trustValue.(bool)
+ return ok && trustForwardHeaders
+}
diff --git a/modules/security/http_filters/security_test.go b/modules/security/http_filters/security_test.go
new file mode 100644
index 000000000..775463a43
--- /dev/null
+++ b/modules/security/http_filters/security_test.go
@@ -0,0 +1,128 @@
+/* Copyright © INFINI LTD. All rights reserved.
+ * Web: https://infinilabs.com
+ * Email: hello#infini.ltd */
+
+package http_filters
+
+import (
+ "net/http"
+ "net/http/httptest"
+ "testing"
+
+ "infini.sh/framework/core/api"
+ httprouter "infini.sh/framework/core/api/router"
+ replaysecurity "infini.sh/framework/core/security/replay"
+)
+
+// Secure-transport enforcement should stop the request before the wrapped UI handler runs.
+func TestSecurityFilterSecureTransportFeature(t *testing.T) {
+ filter := &SecurityFilter{}
+ options := &api.HandlerOptions{}
+ api.SecureTransportOption()(options)
+
+ called := false
+ protected := filter.ApplyFilter(http.MethodPost, "/account/login", options, func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ called = true
+ w.WriteHeader(http.StatusOK)
+ })
+
+ req := httptest.NewRequest(http.MethodPost, "http://console.local/account/login", nil)
+ resp := httptest.NewRecorder()
+ protected(resp, req, nil)
+
+ if called {
+ t.Fatal("expected insecure request to be blocked")
+ }
+ if resp.Code != http.StatusUpgradeRequired {
+ t.Fatalf("expected status %d, got %d", http.StatusUpgradeRequired, resp.Code)
+ }
+}
+
+// When a nonce matches the request scope, the filter should behave like a no-op wrapper.
+func TestSecurityFilterReplayProtectionFeature(t *testing.T) {
+ filter := &SecurityFilter{}
+ options := &api.HandlerOptions{}
+ api.ReplayProtectionOption()(options)
+
+ req := httptest.NewRequest(http.MethodPost, "https://console.local/account/login", nil)
+ nonce, _, err := replaysecurity.IssueReplayNonce(req, http.MethodPost, "/account/login")
+ if err != nil {
+ t.Fatalf("issue replay nonce: %v", err)
+ }
+ req.Header.Set(replaysecurity.HeaderName, nonce)
+
+ called := false
+ protected := filter.ApplyFilter(http.MethodPost, "/account/login", options, func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ called = true
+ w.WriteHeader(http.StatusOK)
+ })
+
+ resp := httptest.NewRecorder()
+ protected(resp, req, nil)
+
+ if !called {
+ t.Fatal("expected replay-protected handler to run")
+ }
+ if resp.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, resp.Code)
+ }
+}
+
+// Missing nonce headers must block replay-protected routes before business logic executes.
+func TestSecurityFilterReplayProtectionRejectsMissingNonce(t *testing.T) {
+ filter := &SecurityFilter{}
+ options := &api.HandlerOptions{}
+ api.ReplayProtectionOption()(options)
+
+ called := false
+ protected := filter.ApplyFilter(http.MethodPost, "/account/login", options, func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ called = true
+ w.WriteHeader(http.StatusOK)
+ })
+
+ req := httptest.NewRequest(http.MethodPost, "https://console.local/account/login", nil)
+ resp := httptest.NewRecorder()
+ protected(resp, req, nil)
+
+ if called {
+ t.Fatal("expected missing nonce to block handler execution")
+ }
+ if resp.Code != http.StatusUnauthorized {
+ t.Fatalf("expected status %d, got %d", http.StatusUnauthorized, resp.Code)
+ }
+}
+
+// Trusted forward headers let deployments behind HTTPS reverse proxies pass transport checks.
+func TestSecurityFilterWithTrustedForwardHeaders(t *testing.T) {
+ filter := &SecurityFilter{}
+ options := &api.HandlerOptions{}
+ api.SecureTransportOption(api.SecureTransportOptions{TrustForwardHeaders: true})(options)
+
+ called := false
+ protected := filter.ApplyFilter(http.MethodPost, "/account/login", options, func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ called = true
+ w.WriteHeader(http.StatusOK)
+ })
+
+ req := httptest.NewRequest(http.MethodPost, "http://console.local/account/login", nil)
+ req.Header.Set("X-Forwarded-Proto", "https")
+ resp := httptest.NewRecorder()
+ protected(resp, req, nil)
+
+ if !called {
+ t.Fatal("expected trusted forwarded proto request to be allowed")
+ }
+ if resp.Code != http.StatusOK {
+ t.Fatalf("expected status %d, got %d", http.StatusOK, resp.Code)
+ }
+}
+
+// Routes that do not opt into trusted proxy headers should stay conservative by default.
+func TestTrustForwardHeadersFromOptionsDefaultsFalse(t *testing.T) {
+ if trustForwardHeadersFromOptions(nil) {
+ t.Fatal("expected nil options to disable trusted forward headers")
+ }
+ if trustForwardHeadersFromOptions(&api.HandlerOptions{}) {
+ t.Fatal("expected missing label to disable trusted forward headers")
+ }
+}
diff --git a/modules/security/native/account_login.go b/modules/security/native/account_login.go
new file mode 100644
index 000000000..6eec79790
--- /dev/null
+++ b/modules/security/native/account_login.go
@@ -0,0 +1,335 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package native
+
+import (
+ "errors"
+ "net/http"
+ "strings"
+ "time"
+
+ log "github.com/cihub/seelog"
+ "infini.sh/framework/core/api"
+ httprouter "infini.sh/framework/core/api/router"
+ "infini.sh/framework/core/orm"
+ "infini.sh/framework/core/security"
+ replaysecurity "infini.sh/framework/core/security/replay"
+ "infini.sh/framework/core/util"
+)
+
+var defaultPasswordChallengeUpgradePersister = func(ctx *orm.Context, user *security.UserAccount) error {
+ return orm.Save(ctx, user)
+}
+
+var persistPasswordChallengeUpgrade = defaultPasswordChallengeUpgradePersister
+
+// RegisterPasswordChallengeUpgradePersister allows applications to override where
+// challenge credentials are persisted after a successful legacy password login.
+func RegisterPasswordChallengeUpgradePersister(persister func(ctx *orm.Context, user *security.UserAccount) error) {
+ if persister == nil {
+ persistPasswordChallengeUpgrade = defaultPasswordChallengeUpgradePersister
+ return
+ }
+ persistPasswordChallengeUpgrade = persister
+}
+
+var (
+ // Keep the password and challenge paths aligned on one user-facing failure message.
+ errInvalidLoginCredentials = errors.New("invalid login or password")
+ // A challenge login must send both the one-time challenge id and the derived proof.
+ errIncompleteChallenge = errors.New("challenge response is incomplete")
+ // Password login keeps requiring the legacy password field when no challenge proof is supplied.
+ errMissingPassword = errors.New("password is required")
+)
+
+func shouldCollapseLoginError(err error) bool {
+ if err == nil {
+ return false
+ }
+
+ switch strings.ToLower(strings.TrimSpace(err.Error())) {
+ case "user not found":
+ return true
+ default:
+ return false
+ }
+}
+
+// accountLoginRequest accepts both the framework-native "login" field and the aliases
+// already used by existing clients while challenge login is rolled out incrementally.
+type accountLoginRequest struct {
+ Login string `json:"login"`
+ Email string `json:"email"`
+ Username string `json:"username"`
+ UserName string `json:"userName"`
+ Password string `json:"password"`
+ ChallengeID string `json:"challenge_id"`
+ Proof string `json:"proof"`
+}
+
+// IssueReplayNonce mints a short-lived nonce bound to the caller and target request scope.
+func IssueReplayNonce(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ var req struct {
+ Method string `json:"method"`
+ Path string `json:"path"`
+ }
+
+ if err := api.DecodeJSON(r, &req); err != nil {
+ api.WriteError(w, err.Error(), http.StatusBadRequest)
+ return
+ }
+
+ nonce, ttl, err := replaysecurity.IssueReplayNonce(r, req.Method, req.Path)
+ if err != nil {
+ api.WriteError(w, err.Error(), http.StatusBadRequest)
+ return
+ }
+
+ api.WriteOKJSON(w, util.MapStr{
+ "status": "ok",
+ "nonce": nonce,
+ "expire_in_seconds": int(ttl / time.Second),
+ })
+}
+
+// LoginChallenge tells the client whether this account can use challenge login and, if so,
+// returns the one-time challenge payload required to derive the proof locally.
+func LoginChallenge(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ var req accountLoginRequest
+ if err := api.DecodeJSON(r, &req); err != nil {
+ api.WriteError(w, err.Error(), http.StatusBadRequest)
+ return
+ }
+
+ login := req.NormalizedLogin()
+ if login == "" {
+ api.WriteError(w, "login is required", http.StatusBadRequest)
+ return
+ }
+
+ exists, user, err := security.GetUserByLogin(login)
+ if err != nil {
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ return
+ }
+
+ api.WriteOKJSON(w, buildLoginChallengeResponse(login, exists, user))
+}
+
+// Login accepts either the legacy password payload or the new challenge proof and then
+// reuses the existing session/token issuance path once the credentials are verified.
+func Login(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
+ var req accountLoginRequest
+ if err := api.DecodeJSON(r, &req); err != nil {
+ api.WriteError(w, err.Error(), http.StatusBadRequest)
+ return
+ }
+
+ login := req.NormalizedLogin()
+ if login == "" {
+ api.WriteError(w, "login is required", http.StatusBadRequest)
+ return
+ }
+
+ usedChallenge := req.ChallengeID != "" || req.Proof != ""
+ exists, user, err := security.GetUserByLogin(login)
+ if err != nil {
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ return
+ }
+ if usedChallenge && (!exists || user == nil) {
+ api.WriteError(w, errInvalidLoginCredentials.Error(), http.StatusForbidden)
+ return
+ }
+
+ if err := validateReplayNonce(r, usedChallenge); err != nil {
+ api.WriteError(w, err.Error(), http.StatusUnauthorized)
+ return
+ }
+
+ usedChallenge, sessionUser, nativeUser, err := authenticateLogin(user, login, req.Password, req.ChallengeID, req.Proof)
+ if err != nil {
+ statusCode := http.StatusForbidden
+ if errors.Is(err, errIncompleteChallenge) || errors.Is(err, errMissingPassword) {
+ statusCode = http.StatusBadRequest
+ }
+ api.WriteError(w, err.Error(), statusCode)
+ return
+ }
+
+ if !usedChallenge && nativeUser != nil {
+ upgradePasswordChallenge(nativeUser, login, req.Password)
+ }
+
+ if err, token := security.AddUserToSession(w, r, sessionUser); err != nil {
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ } else {
+ security.DecorateSessionTokenResponse(token, sessionUser)
+ api.WriteOKJSON(w, token)
+ }
+}
+
+// NormalizedLogin resolves the various historical request field names into one lookup key.
+func (req accountLoginRequest) NormalizedLogin() string {
+ for _, candidate := range []string{req.Login, req.Email, req.Username, req.UserName} {
+ if value := strings.TrimSpace(candidate); value != "" {
+ return value
+ }
+ }
+ return ""
+}
+
+// buildLoginChallengeResponse returns plain login for existing legacy accounts that have
+// not been upgraded with challenge material yet. Accounts that do not exist still receive
+// a fake challenge payload to avoid user enumeration.
+func buildLoginChallengeResponse(login string, exists bool, user *security.UserAccount) util.MapStr {
+ if exists && !security.CanUsePasswordChallenge(user) {
+ return util.MapStr{
+ "status": "ok",
+ "method": "plain",
+ }
+ }
+
+ salt := util.GenerateSecureString(32)
+ if exists && security.CanUsePasswordChallenge(user) {
+ salt = user.PasswordSalt
+ }
+
+ // The challenge payload gives clients everything needed to derive a proof
+ // locally without sending the raw password back to the server.
+ challenge := security.NewLoginChallenge(login)
+ return util.MapStr{
+ "status": "ok",
+ "method": security.PasswordChallengeMethod,
+ "algorithm": security.PasswordChallengeAlgorithm,
+ "iterations": security.PasswordChallengeIterations,
+ "challenge_id": challenge.ID,
+ "nonce": challenge.Nonce,
+ "salt": salt,
+ }
+}
+
+// authenticateLogin selects the correct credential validation path based on the request body.
+func authenticateLogin(user *security.UserAccount, login, password, challengeID, proof string) (bool, *security.UserSessionInfo, *security.UserAccount, error) {
+ if challengeID != "" || proof != "" {
+ if challengeID == "" || proof == "" {
+ return true, nil, nil, errIncompleteChallenge
+ }
+
+ if user == nil {
+ return true, nil, nil, errInvalidLoginCredentials
+ }
+ challenge, err := security.ConsumeLoginChallenge(challengeID, login)
+ if err != nil || !security.CanUsePasswordChallenge(user) {
+ return true, nil, nil, errInvalidLoginCredentials
+ }
+ if !security.VerifyPasswordProof(user.PasswordVerifier, login, challenge.ID, challenge.Nonce, proof) {
+ return true, nil, nil, errInvalidLoginCredentials
+ }
+ return true, newNativeSession(user, login), user, nil
+ }
+
+ if password == "" {
+ return false, nil, nil, errMissingPassword
+ }
+
+ if user != nil {
+ if err := security.VerifyPassword(user, password); err == nil {
+ return false, newNativeSession(user, login), user, nil
+ }
+ }
+
+ sessionUser, err := security.AuthenticateAccountPasswordLogin(login, password)
+ if err != nil {
+ if shouldCollapseLoginError(err) {
+ return false, nil, nil, errInvalidLoginCredentials
+ }
+ return false, nil, nil, err
+ }
+ if sessionUser != nil {
+ return false, sessionUser, nil, nil
+ }
+
+ return false, nil, nil, errInvalidLoginCredentials
+}
+
+// validateReplayNonce keeps challenge login replay-safe while leaving older password-only
+// clients working until they adopt the explicit nonce negotiation endpoint.
+func validateReplayNonce(r *http.Request, required bool) error {
+ nonce := strings.TrimSpace(r.Header.Get(replaysecurity.HeaderName))
+ if nonce == "" && !required {
+ // Keep the original password login path backward compatible: upgraded clients
+ // send replay nonces, while older clients can still post passwords directly.
+ return nil
+ }
+ return replaysecurity.ValidateAndConsumeReplayNonce(r)
+}
+
+// upgradePasswordChallenge backfills verifier material after a successful legacy login so
+// existing native accounts can move onto the challenge flow without an offline migration.
+func upgradePasswordChallenge(user *security.UserAccount, login, password string) {
+ if user == nil || password == "" || security.CanUsePasswordChallenge(user) {
+ return
+ }
+
+ if err := security.EnsurePasswordChallenge(user, password); err != nil {
+ log.Warnf("failed to derive password challenge for user [%s]: %v", user.Email, err)
+ return
+ }
+
+ // Persist the verifier after a successful legacy password login so subsequent
+ // logins can move onto the challenge flow without an explicit migration step.
+ // This upgrade is best-effort; the current login already succeeded, so it should
+ // not wait for an index refresh before returning to the caller.
+ if user.ID == "" {
+ userLogin := strings.TrimSpace(user.Email)
+ if userLogin == "" {
+ userLogin = strings.TrimSpace(login)
+ }
+ if userLogin != "" {
+ user.ID = getUIDByEmail(userLogin)
+ }
+ }
+ ctx := orm.NewContext()
+ ctx.DirectAccess()
+ if err := persistPasswordChallengeUpgrade(ctx, user); err != nil {
+ log.Warnf("failed to persist password challenge for user [%s]: %v", user.Email, err)
+ }
+}
+
+// newNativeSession converts a native account record into the existing framework session claims.
+func newNativeSession(user *security.UserAccount, login string) *security.UserSessionInfo {
+ userLogin := strings.TrimSpace(user.Email)
+ if userLogin == "" {
+ userLogin = login
+ }
+
+ session := &security.UserSessionInfo{
+ Provider: security.DefaultNativeAuthBackend,
+ Login: userLogin,
+ Roles: append([]string(nil), user.Roles...),
+ }
+ session.SetUserID(user.ID)
+ return session
+}
diff --git a/modules/security/native/account_login_test.go b/modules/security/native/account_login_test.go
new file mode 100644
index 000000000..9f64ec30a
--- /dev/null
+++ b/modules/security/native/account_login_test.go
@@ -0,0 +1,416 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package native
+
+import (
+ "bytes"
+ "encoding/json"
+ "errors"
+ "net/http"
+ "net/http/httptest"
+ "testing"
+ "time"
+
+ "golang.org/x/crypto/bcrypt"
+ "infini.sh/framework/core/orm"
+ "infini.sh/framework/core/security"
+ replaysecurity "infini.sh/framework/core/security/replay"
+)
+
+type testAccountPasswordLoginProvider struct{}
+type testMissingUserAccountPasswordLoginProvider struct{}
+
+type testChallengeAuthenticationBackend struct{}
+
+func (testChallengeAuthenticationBackend) GetUserByID(id string) (bool, *security.UserAccount, error) {
+ return false, nil, nil
+}
+
+func (testChallengeAuthenticationBackend) GetUserByLogin(login string) (bool, *security.UserAccount, error) {
+ if login != "bridge-admin" {
+ return false, nil, nil
+ }
+ user := &security.UserAccount{Email: "bridge-admin"}
+ if err := security.SetPassword(user, "StrongPassw0rd!"); err != nil {
+ return false, nil, err
+ }
+ user.ID = "bridge-admin-id"
+ return true, user, nil
+}
+
+func (testChallengeAuthenticationBackend) CreateUser(name, login, password string, force bool) (*security.UserAccount, error) {
+ return nil, nil
+}
+
+func (testAccountPasswordLoginProvider) AuthenticateByPassword(login, password string) (*security.UserSessionInfo, error) {
+ if login != "ldap-user" || password != "StrongPassw0rd!" {
+ return nil, nil
+ }
+
+ sessionUser := &security.UserSessionInfo{
+ Provider: "ldap",
+ Login: login,
+ Roles: []string{"viewer"},
+ }
+ sessionUser.SetUserID("ldap-user-id")
+ return sessionUser, nil
+}
+
+func (testMissingUserAccountPasswordLoginProvider) AuthenticateByPassword(login, password string) (*security.UserSessionInfo, error) {
+ if login != "missing-user" {
+ return nil, nil
+ }
+
+ return nil, errors.New("user not found")
+}
+
+// The request payload accepts multiple historical login field names during rollout.
+func TestAccountLoginRequestNormalizedLogin(t *testing.T) {
+ req := accountLoginRequest{
+ Email: "admin@example.org",
+ Username: "ignored@example.org",
+ }
+
+ if got := req.NormalizedLogin(); got != "admin@example.org" {
+ t.Fatalf("expected email to be preferred, got %q", got)
+ }
+}
+
+// Password login remains the backward-compatible path for accounts and clients not yet upgraded.
+func TestAuthenticateLoginWithPassword(t *testing.T) {
+ user := &security.UserAccount{Email: "admin@example.org"}
+ if err := security.SetPassword(user, "StrongPassw0rd!"); err != nil {
+ t.Fatalf("set password: %v", err)
+ }
+
+ usedChallenge, sessionUser, nativeUser, err := authenticateLogin(user, user.Email, "StrongPassw0rd!", "", "")
+ if err != nil {
+ t.Fatalf("authenticate login: %v", err)
+ }
+ if usedChallenge {
+ t.Fatal("expected password login path")
+ }
+ if sessionUser == nil || nativeUser == nil {
+ t.Fatalf("expected native password login state, got session=%#v native=%#v", sessionUser, nativeUser)
+ }
+}
+
+// Challenge login should succeed once the account already has verifier material.
+func TestAuthenticateLoginWithChallenge(t *testing.T) {
+ user := &security.UserAccount{Email: "admin@example.org"}
+ if err := security.SetPassword(user, "StrongPassw0rd!"); err != nil {
+ t.Fatalf("set password: %v", err)
+ }
+
+ challenge := security.NewLoginChallenge(user.Email)
+ proof, err := security.BuildPasswordProof(user.PasswordVerifier, user.Email, challenge.ID, challenge.Nonce)
+ if err != nil {
+ t.Fatalf("build password proof: %v", err)
+ }
+
+ usedChallenge, sessionUser, nativeUser, err := authenticateLogin(user, user.Email, "", challenge.ID, proof)
+ if err != nil {
+ t.Fatalf("authenticate login: %v", err)
+ }
+ if !usedChallenge {
+ t.Fatal("expected challenge login path")
+ }
+ if sessionUser == nil || nativeUser == nil {
+ t.Fatalf("expected native challenge login state, got session=%#v native=%#v", sessionUser, nativeUser)
+ }
+}
+
+// Partially supplied challenge payloads should fail distinctly from bad credentials.
+func TestAuthenticateLoginRejectsIncompleteChallenge(t *testing.T) {
+ user := &security.UserAccount{Email: "admin@example.org"}
+ if err := security.SetPassword(user, "StrongPassw0rd!"); err != nil {
+ t.Fatalf("set password: %v", err)
+ }
+
+ _, _, _, err := authenticateLogin(user, user.Email, "", "challenge-id", "")
+ if !errors.Is(err, errIncompleteChallenge) {
+ t.Fatalf("expected incomplete challenge error, got %v", err)
+ }
+}
+
+// Incorrect proofs should collapse to the same user-facing error as bad passwords.
+func TestAuthenticateLoginRejectsWrongProof(t *testing.T) {
+ user := &security.UserAccount{Email: "admin@example.org"}
+ if err := security.SetPassword(user, "StrongPassw0rd!"); err != nil {
+ t.Fatalf("set password: %v", err)
+ }
+
+ challenge := security.NewLoginChallenge(user.Email)
+ _, _, _, err := authenticateLogin(user, user.Email, "", challenge.ID, "bad-proof")
+ if !errors.Is(err, errInvalidLoginCredentials) {
+ t.Fatalf("expected invalid credential error, got %v", err)
+ }
+}
+
+// Applications can attach non-native password realms to the shared framework login flow.
+func TestAuthenticateLoginFallsBackToRegisteredPasswordProvider(t *testing.T) {
+ security.RegisterAccountPasswordLoginProvider("test-account-login", testAccountPasswordLoginProvider{})
+
+ usedChallenge, sessionUser, nativeUser, err := authenticateLogin(nil, "ldap-user", "StrongPassw0rd!", "", "")
+ if err != nil {
+ t.Fatalf("authenticate login: %v", err)
+ }
+ if usedChallenge {
+ t.Fatal("expected password fallback path")
+ }
+ if nativeUser != nil {
+ t.Fatalf("expected no native user for fallback path, got %#v", nativeUser)
+ }
+ if sessionUser == nil || sessionUser.Provider != "ldap" {
+ t.Fatalf("expected ldap session user, got %#v", sessionUser)
+ }
+}
+
+func TestAuthenticateLoginCollapsesMissingUserProviderError(t *testing.T) {
+ security.RegisterAccountPasswordLoginProvider("test-account-login-missing-user", testMissingUserAccountPasswordLoginProvider{})
+
+ _, _, _, err := authenticateLogin(nil, "missing-user", "StrongPassw0rd!", "", "")
+ if !errors.Is(err, errInvalidLoginCredentials) {
+ t.Fatalf("expected invalid credential error for missing user, got %v", err)
+ }
+}
+
+// Older accounts intentionally advertise plain login until their verifier is available.
+func TestBuildLoginChallengeResponseFallsBackToPlain(t *testing.T) {
+ user := &security.UserAccount{Email: "admin@example.org"}
+ resp := buildLoginChallengeResponse(user.Email, true, user)
+
+ if got := resp["method"]; got != "plain" {
+ t.Fatalf("expected plain fallback, got %v", got)
+ }
+ if _, ok := resp["challenge_id"]; ok {
+ t.Fatal("did not expect challenge payload for plain fallback")
+ }
+}
+
+func TestBuildLoginChallengeResponseFakesChallengeForMissingUser(t *testing.T) {
+ resp := buildLoginChallengeResponse("missing@example.org", false, nil)
+
+ if got := resp["method"]; got != security.PasswordChallengeMethod {
+ t.Fatalf("expected fake challenge method for missing user, got %v", got)
+ }
+ if resp["challenge_id"] == "" {
+ t.Fatal("expected fake challenge id for missing user")
+ }
+ if resp["nonce"] == "" {
+ t.Fatal("expected fake nonce for missing user")
+ }
+ if resp["salt"] == "" {
+ t.Fatal("expected fake salt for missing user")
+ }
+}
+
+// Upgraded accounts should return the exact challenge inputs the client needs next.
+func TestBuildLoginChallengeResponseReturnsChallenge(t *testing.T) {
+ user := &security.UserAccount{Email: "admin@example.org"}
+ if err := security.SetPassword(user, "StrongPassw0rd!"); err != nil {
+ t.Fatalf("set password: %v", err)
+ }
+
+ resp := buildLoginChallengeResponse(user.Email, true, user)
+ if got := resp["method"]; got != security.PasswordChallengeMethod {
+ t.Fatalf("expected challenge method, got %v", got)
+ }
+ if resp["challenge_id"] == "" {
+ t.Fatal("expected challenge id to be returned")
+ }
+ if resp["nonce"] == "" {
+ t.Fatal("expected nonce to be returned")
+ }
+ if resp["salt"] != user.PasswordSalt {
+ t.Fatal("expected challenge response to expose password salt")
+ }
+}
+
+func TestLoginChallengeUsesRegisteredAuthenticationBackend(t *testing.T) {
+ security.RegisterAuthenticationProvider("test-login-challenge-provider", testChallengeAuthenticationBackend{})
+
+ body := bytes.NewBufferString(`{"login":"bridge-admin"}`)
+ req := httptest.NewRequest(http.MethodPost, "/account/login/challenge", body)
+ req.Header.Set("Content-Type", "application/json")
+ recorder := httptest.NewRecorder()
+
+ LoginChallenge(recorder, req, nil)
+
+ if recorder.Code != http.StatusOK {
+ t.Fatalf("expected 200 response, got %d: %s", recorder.Code, recorder.Body.String())
+ }
+
+ var resp map[string]interface{}
+ if err := json.Unmarshal(recorder.Body.Bytes(), &resp); err != nil {
+ t.Fatalf("decode response: %v", err)
+ }
+ if got := resp["method"]; got != security.PasswordChallengeMethod {
+ t.Fatalf("expected challenge method from registered provider, got %v", got)
+ }
+ if resp["challenge_id"] == "" {
+ t.Fatal("expected challenge id from registered provider")
+ }
+}
+
+// Legacy password clients keep working even before they learn the replay-nonce preflight.
+func TestValidateReplayNonceAllowsLegacyPasswordLoginWithoutNonce(t *testing.T) {
+ req := httptest.NewRequest(http.MethodPost, "/account/login", nil)
+ if err := validateReplayNonce(req, false); err != nil {
+ t.Fatalf("expected missing nonce to be allowed for legacy password login, got %v", err)
+ }
+}
+
+// Challenge logins must enforce nonce usage immediately because the frontend already negotiated it.
+func TestValidateReplayNonceRequiresNonceForChallengeLogin(t *testing.T) {
+ req := httptest.NewRequest(http.MethodPost, "/account/login", nil)
+ if err := validateReplayNonce(req, true); err == nil {
+ t.Fatal("expected missing nonce to be rejected for challenge login")
+ }
+}
+
+// Once a nonce is explicitly issued for /account/login it should validate on that exact route.
+func TestValidateReplayNonceConsumesIssuedNonce(t *testing.T) {
+ req := httptest.NewRequest(http.MethodPost, "/account/login", nil)
+ nonce, _, err := replaysecurity.IssueReplayNonce(req, http.MethodPost, "/account/login")
+ if err != nil {
+ t.Fatalf("issue replay nonce: %v", err)
+ }
+ req.Header.Set(replaysecurity.HeaderName, nonce)
+
+ if err := validateReplayNonce(req, true); err != nil {
+ t.Fatalf("expected issued nonce to validate, got %v", err)
+ }
+}
+
+// Native sessions should still be constructible even when the stored account email is blank.
+func TestNewNativeSessionFallsBackToRequestedLogin(t *testing.T) {
+ user := &security.UserAccount{Email: "", Roles: []string{security.RoleAdmin}}
+ user.ID = "user-1"
+
+ session := newNativeSession(user, "admin@example.org")
+ if session.Login != "admin@example.org" {
+ t.Fatalf("expected requested login fallback, got %q", session.Login)
+ }
+ if session.Provider != security.DefaultNativeAuthBackend {
+ t.Fatalf("expected native provider, got %q", session.Provider)
+ }
+}
+
+func TestUpgradePasswordChallengePersistsLegacyAdminByLogin(t *testing.T) {
+ originalPersist := persistPasswordChallengeUpgrade
+ defer func() {
+ persistPasswordChallengeUpgrade = originalPersist
+ }()
+
+ var persisted *security.UserAccount
+ persistPasswordChallengeUpgrade = func(ctx *orm.Context, user *security.UserAccount) error {
+ copied := *user
+ persisted = &copied
+ return nil
+ }
+
+ user := &security.UserAccount{Name: "admin"}
+ hash, err := bcrypt.GenerateFromPassword([]byte("StrongPassw0rd!"), bcrypt.DefaultCost)
+ if err != nil {
+ t.Fatalf("generate password hash: %v", err)
+ }
+ user.Password = string(hash)
+
+ upgradePasswordChallenge(user, "admin", "StrongPassw0rd!")
+
+ if persisted == nil {
+ t.Fatal("expected legacy admin upgrade to be persisted")
+ }
+ if persisted.ID != getUIDByEmail("admin") {
+ t.Fatalf("expected fallback id %q, got %q", getUIDByEmail("admin"), persisted.ID)
+ }
+ if persisted.PasswordSalt == "" || persisted.PasswordVerifier == "" {
+ t.Fatal("expected challenge credentials to be populated before persisting")
+ }
+}
+
+func TestUpgradePasswordChallengeSkipsExistingChallengeUser(t *testing.T) {
+ originalPersist := persistPasswordChallengeUpgrade
+ defer func() {
+ persistPasswordChallengeUpgrade = originalPersist
+ }()
+
+ called := false
+ persistPasswordChallengeUpgrade = func(ctx *orm.Context, user *security.UserAccount) error {
+ called = true
+ return nil
+ }
+
+ user := &security.UserAccount{Email: "admin@example.org"}
+ if err := security.SetPassword(user, "StrongPassw0rd!"); err != nil {
+ t.Fatalf("set password: %v", err)
+ }
+
+ upgradePasswordChallenge(user, user.Email, "StrongPassw0rd!")
+ if !security.CanUsePasswordChallenge(user) {
+ t.Fatal("expected challenge material to be available")
+ }
+ called = false
+
+ upgradePasswordChallenge(user, user.Email, "StrongPassw0rd!")
+ if called {
+ t.Fatal("did not expect already-upgraded account to be persisted again")
+ }
+}
+
+// The framework login response keeps the console frontend contract while the handler
+// implementation moves from console into framework-owned routes.
+func TestDecorateLoginResponseAddsConsoleCompatibilityFields(t *testing.T) {
+ session := &security.UserSessionInfo{
+ Provider: security.DefaultNativeAuthBackend,
+ Login: "admin@example.org",
+ Roles: []string{security.RoleAdmin},
+ Permissions: []security.PermissionKey{security.GetSimplePermission("generic", "unit", security.Read)},
+ }
+ session.SetUserID("user-1")
+
+ token := map[string]interface{}{
+ "status": "ok",
+ "expire_in": time.Now().Unix() + 3600,
+ }
+ security.DecorateSessionTokenResponse(token, session)
+
+ if token["username"] != session.Login {
+ t.Fatalf("expected username %q, got %v", session.Login, token["username"])
+ }
+ if token["id"] != session.UserID {
+ t.Fatalf("expected id %q, got %v", session.UserID, token["id"])
+ }
+ if token["expires_at"] == nil {
+ t.Fatal("expected expires_at to be populated")
+ }
+ if expireIn, ok := token["expire_in"].(int64); !ok || expireIn <= 0 || expireIn > 3600 {
+ t.Fatalf("expected expire_in to become remaining lifetime seconds, got %#v", token["expire_in"])
+ }
+ privilege, ok := token["privilege"].([]security.PermissionKey)
+ if !ok || len(privilege) == 0 {
+ t.Fatalf("expected privilege list to be populated, got %#v", token["privilege"])
+ }
+}
diff --git a/modules/security/native/authorization.go b/modules/security/native/authorization.go
index 06042cbb4..73d377b47 100644
--- a/modules/security/native/authorization.go
+++ b/modules/security/native/authorization.go
@@ -8,6 +8,7 @@ import (
"context"
"net/http"
+ log "github.com/cihub/seelog"
"infini.sh/framework/core/api"
httprouter "infini.sh/framework/core/api/router"
"infini.sh/framework/core/elastic"
@@ -17,6 +18,15 @@ import (
"infini.sh/framework/core/util"
)
+const (
+ errInvalidCurrentUser = "invalid user"
+ errInvalidRole = "invalid role"
+ errCannotUpdateOwnRole = "you can not update the roles for you"
+ errReservedRoleName = "can not use the reserved role name"
+ errRoleAlreadyExists = "same role name already exists"
+ errRoleAssignedToUsers = "role is still assigned to users"
+)
+
func GetRole(w http.ResponseWriter, req *http.Request, ps httprouter.Params) {
id := ps.MustGetParameter("id")
@@ -28,8 +38,12 @@ func GetRole(w http.ResponseWriter, req *http.Request, ps httprouter.Params) {
ctx.PermissionScope(security.PermissionScopePlatform)
exists, err := orm.GetV2(ctx, &obj)
- if !exists || err != nil {
- api.NotFoundResponse(id)
+ if err != nil {
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ return
+ }
+ if !exists {
+ api.WriteJSON(w, api.NotFoundResponse(id), http.StatusNotFound)
return
}
@@ -44,7 +58,7 @@ func UpdateRole(w http.ResponseWriter, req *http.Request, ps httprouter.Params)
obj := security.UserRole{}
err := api.DecodeJSON(req, &obj)
if err != nil {
- api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ api.WriteError(w, err.Error(), http.StatusBadRequest)
return
}
@@ -56,16 +70,23 @@ func UpdateRole(w http.ResponseWriter, req *http.Request, ps httprouter.Params)
userID := sessionUser.MustGetUserID()
_, account, err := security.GetUserByID(userID)
- if account == nil || err != nil {
- panic("invalid user")
+ if err != nil {
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ return
+ }
+ if account == nil {
+ api.WriteError(w, errInvalidCurrentUser, http.StatusUnauthorized)
+ return
}
_, role := GetRoleByID(id)
if role == nil {
- panic("invalid role")
+ api.WriteError(w, errInvalidRole, http.StatusNotFound)
+ return
}
if util.ContainsAnyInArray(role.Name, account.Roles) {
- panic("you can not update the roles for you")
+ api.WriteError(w, errCannotUpdateOwnRole, http.StatusForbidden)
+ return
}
}
@@ -89,8 +110,28 @@ func DeleteRole(w http.ResponseWriter, req *http.Request, ps httprouter.Params)
obj.ID = id
ctx := orm.NewContextWithParent(req.Context())
ctx.DirectAccess()
+
+ exists, err := orm.GetV2(ctx, &obj)
+ if err != nil {
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ return
+ }
+ if !exists {
+ api.WriteJSON(w, api.NotFoundResponse(id), http.StatusNotFound)
+ return
+ }
+ inUse, err := roleHasAssignedUsers(req.Context(), obj.Name)
+ if err != nil {
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ return
+ }
+ if inUse {
+ api.WriteError(w, errRoleAssignedToUsers, http.StatusConflict)
+ return
+ }
+
ctx.Refresh = orm.WaitForRefresh
- err := orm.Delete(ctx, &obj)
+ err = orm.Delete(ctx, &obj)
if err != nil {
api.WriteError(w, err.Error(), http.StatusInternalServerError)
return
@@ -164,19 +205,21 @@ func CreateRole(w http.ResponseWriter, req *http.Request, ps httprouter.Params)
var obj = &security.UserRole{}
err := api.DecodeJSON(req, obj)
if err != nil {
- api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ api.WriteError(w, err.Error(), http.StatusBadRequest)
return
}
if obj.Name == "admin" {
- panic("can not use the reserved role name")
+ api.WriteError(w, errReservedRoleName, http.StatusBadRequest)
+ return
}
api.MustValidateInput(w, obj)
exists, _ := GetRoleByName(obj.Name)
if exists {
- panic("same role name already exists")
+ api.WriteError(w, errRoleAlreadyExists, http.StatusConflict)
+ return
}
ctx := orm.NewContextWithParent(req.Context())
@@ -242,7 +285,8 @@ func (provider *SecurityBackendProvider) GetPermissionKeysByRoles(ctx1 context.C
result := []security.UserRole{}
err, _ := elastic.SearchV2WithResultItemMapper(ctx, &result, qb, nil)
if err != nil {
- panic(err)
+ log.Errorf("failed to load permissions for roles %v: %v", roles, err)
+ return []security.PermissionKey{}
}
allowed := make(map[security.PermissionKey]struct{}, 128)
@@ -269,3 +313,22 @@ func (provider *SecurityBackendProvider) GetPermissionKeysByRoles(ctx1 context.C
}
return keys
}
+
+func roleHasAssignedUsers(ctx1 context.Context, roleName string) (bool, error) {
+ if roleName == "" {
+ return false, nil
+ }
+
+ ctx := orm.NewContextWithParent(ctx1)
+ ctx.DirectReadAccess()
+ ctx.PermissionScope(security.PermissionScopePlatform)
+ orm.WithModel(ctx, &security.UserAccount{})
+
+ qb := orm.NewQuery()
+ qb.Must(orm.TermQuery("roles", roleName))
+ err, result := elastic.SearchV2WithResultItemMapper(ctx, nil, qb, nil)
+ if err != nil {
+ return false, err
+ }
+ return result != nil && result.Total > 0, nil
+}
diff --git a/modules/security/native/entity.go b/modules/security/native/entity.go
index 0153985b6..7b31c2f73 100644
--- a/modules/security/native/entity.go
+++ b/modules/security/native/entity.go
@@ -7,6 +7,7 @@ package native
import (
"context"
+ log "github.com/cihub/seelog"
"infini.sh/framework/core/elastic"
"infini.sh/framework/core/entity_card"
"infini.sh/framework/core/orm"
@@ -47,7 +48,8 @@ func (this *UserEntityProvider) GenEntityLabel(ctx1 context.Context, t string, i
out := []security.UserAccount{}
err, _ := elastic.SearchV2WithResultItemMapper(ctx, &out, builder, nil)
if err != nil {
- panic(err)
+ log.Errorf("failed to load user entity labels for ids %v: %v", ids, err)
+ return output
}
for _, a := range out {
diff --git a/modules/security/native/init.go b/modules/security/native/init.go
index 759719abe..d591dba0d 100644
--- a/modules/security/native/init.go
+++ b/modules/security/native/init.go
@@ -19,6 +19,8 @@ func Init() {
security.RegisterAuthenticationProvider(security.DefaultNativeAuthBackend, &provider)
security.RegisterAuthorizationProvider(security.DefaultNativeAuthBackend, &provider)
+ RegisterPublicUIAuthRoutes()
+
orm.MustRegisterSchemaWithIndexName(&security.UserAccount{}, "app-users")
orm.MustRegisterSchemaWithIndexName(&security.UserRole{}, "app-roles")
@@ -58,3 +60,24 @@ func Init() {
}
}
+
+func RegisterPublicUIAuthRoutes() {
+ secureViaProxy := api.SecureTransportOptions{TrustForwardHeaders: true}
+ api.HandleUIMethod(api.POST, "/account/replay_nonce",
+ api.RequireSecureTransport(IssueReplayNonce, secureViaProxy),
+ api.AllowPublicAccess(),
+ api.AllowOPTIONSS(),
+ api.Feature(api.FeatureCORS))
+
+ api.HandleUIMethod(api.POST, "/account/login/challenge",
+ api.RequireSecureTransport(LoginChallenge, secureViaProxy),
+ api.AllowPublicAccess(),
+ api.AllowOPTIONSS(),
+ api.Feature(api.FeatureCORS))
+
+ api.HandleUIMethod(api.POST, "/account/login",
+ api.RequireSecureTransport(Login, secureViaProxy),
+ api.AllowPublicAccess(),
+ api.AllowOPTIONSS(),
+ api.Feature(api.FeatureCORS))
+}
diff --git a/modules/security/native/principal.go b/modules/security/native/principal.go
index f1e52c594..6dcc15710 100644
--- a/modules/security/native/principal.go
+++ b/modules/security/native/principal.go
@@ -18,7 +18,8 @@ func SearchPrincipals(w http.ResponseWriter, req *http.Request, ps httprouter.Pa
builder, err := orm.NewQueryBuilderFromRequest(req, "id", "name", "email")
if err != nil {
- panic(err)
+ api.WriteError(w, err.Error(), http.StatusBadRequest)
+ return
}
ctx := orm.NewContextWithParent(req.Context())
ctx.DirectReadAccess()
@@ -26,7 +27,8 @@ func SearchPrincipals(w http.ResponseWriter, req *http.Request, ps httprouter.Pa
out := []security.UserAccount{}
err, res := elastic.SearchV2WithResultItemMapper(ctx, &out, builder, nil)
if err != nil {
- panic(err)
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ return
}
// use the generic type correctly
diff --git a/modules/security/native/user.go b/modules/security/native/user.go
index 17b71c8aa..99186f6d2 100644
--- a/modules/security/native/user.go
+++ b/modules/security/native/user.go
@@ -5,19 +5,27 @@
package native
import (
+ "fmt"
"net/http"
log "github.com/cihub/seelog"
- "golang.org/x/crypto/bcrypt"
"infini.sh/framework/core/api"
httprouter "infini.sh/framework/core/api/router"
"infini.sh/framework/core/elastic"
+ cerr "infini.sh/framework/core/errors"
"infini.sh/framework/core/global"
"infini.sh/framework/core/orm"
"infini.sh/framework/core/security"
"infini.sh/framework/core/util"
)
+const (
+ errCannotUpdateOwnRoles = "sorry, you can not update your roles"
+ errCannotDeleteSelf = "you can not delete yourself"
+ errInsecurePassword = "password does not meet security requirements"
+ errEmailAlreadyExists = "email already existed"
+)
+
func GetUser(w http.ResponseWriter, req *http.Request, ps httprouter.Params) {
id := ps.MustGetParameter("id")
@@ -26,10 +34,11 @@ func GetUser(w http.ResponseWriter, req *http.Request, ps httprouter.Params) {
ctx := orm.NewContextWithParent(req.Context())
exists, err := orm.GetV2(ctx, &obj)
if err != nil {
- panic(err)
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ return
}
if !exists {
- api.NotFoundResponse(id)
+ api.WriteJSON(w, api.NotFoundResponse(id), http.StatusNotFound)
return
}
@@ -43,7 +52,8 @@ func UpdateUser(w http.ResponseWriter, req *http.Request, ps httprouter.Params)
obj := security.UserAccount{}
err := api.DecodeJSON(req, &obj)
if err != nil {
- panic(err)
+ api.WriteError(w, err.Error(), http.StatusBadRequest)
+ return
}
api.MustValidateInput(w, obj)
@@ -52,10 +62,11 @@ func UpdateUser(w http.ResponseWriter, req *http.Request, ps httprouter.Params)
oldObj.ID = id
exists, err := orm.GetV2(ctx, &oldObj)
if err != nil {
- panic(err)
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ return
}
if !exists {
- api.NotFoundResponse(id)
+ api.WriteJSON(w, api.NotFoundResponse(id), http.StatusNotFound)
return
}
@@ -68,26 +79,32 @@ func UpdateUser(w http.ResponseWriter, req *http.Request, ps httprouter.Params)
if userID == id {
//user can't update self's role
if !util.CompareStringArray(obj.Roles, oldObj.Roles) {
- panic("sorry, you can not update your roles")
+ api.WriteError(w, errCannotUpdateOwnRoles, http.StatusForbidden)
+ return
}
}
if obj.Password == "" {
+ // Preserve the verifier material on metadata-only updates so editing roles,
+ // names, or other fields does not silently disable challenge login.
obj.Password = oldObj.Password
+ obj.PasswordSalt = oldObj.PasswordSalt
+ obj.PasswordVerifier = oldObj.PasswordVerifier
} else {
- if !util.ValidateSecure(obj.Password) {
- panic("should be secured password")
+ if err := validateSecurePassword(obj.Password); err != nil {
+ api.WriteError(w, err.Error(), http.StatusBadRequest)
+ return
}
- hash, err := bcrypt.GenerateFromPassword([]byte(obj.Password), bcrypt.DefaultCost)
- if err != nil {
- panic(err)
+ if err := security.SetPassword(&obj, obj.Password); err != nil {
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ return
}
- obj.Password = string(hash)
}
ctx.Refresh = orm.WaitForRefresh
err = orm.Update(ctx, &obj)
if err != nil {
- panic(err)
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ return
}
security.IncreasePermissionVersion()
@@ -103,13 +120,15 @@ func DeleteUser(w http.ResponseWriter, req *http.Request, ps httprouter.Params)
sessionUser := security.MustGetUserFromContext(ctx)
userID := sessionUser.MustGetUserID()
if userID == id {
- panic("you can not delete yourself")
+ api.WriteError(w, errCannotDeleteSelf, http.StatusForbidden)
+ return
}
ctx.Refresh = orm.WaitForRefresh
err := orm.Delete(ctx, &obj)
if err != nil {
- panic(err)
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ return
}
api.WriteDeletedOKJSON(w, obj.ID)
@@ -118,7 +137,8 @@ func DeleteUser(w http.ResponseWriter, req *http.Request, ps httprouter.Params)
func SearchUser(w http.ResponseWriter, req *http.Request, ps httprouter.Params) {
builder, err := orm.NewQueryBuilderFromRequest(req, "id", "name", "email")
if err != nil {
- panic(err)
+ api.WriteError(w, err.Error(), http.StatusBadRequest)
+ return
}
ctx := orm.NewContextWithParent(req.Context())
ctx.DirectReadAccess()
@@ -128,12 +148,13 @@ func SearchUser(w http.ResponseWriter, req *http.Request, ps httprouter.Params)
orm.WithModel(ctx, &security.UserAccount{})
res, err := orm.SearchV2(ctx, builder)
if err != nil {
- panic(err)
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ return
}
_, err = api.Write(w, res.Payload.([]byte))
if err != nil {
- panic(err)
+ api.Error(w, err)
}
}
@@ -152,14 +173,8 @@ func GetUserByLogin(email string) (bool, *security.UserAccount, error) {
if err != nil {
return false, nil, err
}
- if len(items) > 0 {
- if len(items) == 1 {
- return true, &items[0], nil
- } else {
- log.Warnf("invalid users, more than one account was associated with the same email: %v", email)
- }
- }
- return false, nil, nil
+
+ return resolveUserByLogin(email, items)
}
func (provider *SecurityBackendProvider) GetUserByLogin(email string) (bool, *security.UserAccount, error) {
@@ -186,17 +201,17 @@ func (provider *SecurityBackendProvider) GetUserByID(id string) (bool, *security
func (provider *SecurityBackendProvider) CreateUser(name, email, password string, force bool) (*security.UserAccount, error) {
- if !util.ValidateSecure(password) {
- panic("should be secured password")
+ if err := validateSecurePassword(password); err != nil {
+ return nil, err
}
exists, account, err := GetUserByLogin(email)
if err != nil {
- panic(err)
+ return nil, err
}
if exists && !force {
- panic("email already existed")
+ return nil, cerr.NewWithHTTPCode(http.StatusConflict, errEmailAlreadyExists)
}
var obj = &security.UserAccount{}
@@ -204,24 +219,22 @@ func (provider *SecurityBackendProvider) CreateUser(name, email, password string
log.Warn("email already exists, will be replaced")
obj.ID = account.ID
} else {
- obj.ID = getUIDByEmail(obj.Email)
+ obj.ID = getUIDByEmail(email)
}
- hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
- if err != nil {
- panic(err)
- }
obj.Name = name
obj.Email = email
obj.Roles = []string{security.RoleAdmin}
- obj.Password = string(hash)
+ if err := security.SetPassword(obj, password); err != nil {
+ return nil, err
+ }
ctx := orm.NewContext()
ctx.DirectAccess()
ctx.Refresh = orm.WaitForRefresh
err = orm.Save(ctx, obj)
if err != nil {
- panic(err)
+ return nil, err
}
return obj, nil
}
@@ -234,38 +247,62 @@ func CreateUser(w http.ResponseWriter, req *http.Request, ps httprouter.Params)
var obj = &security.UserAccount{}
err := api.DecodeJSON(req, obj)
if err != nil {
- panic(err)
+ api.WriteError(w, err.Error(), http.StatusBadRequest)
+ return
}
api.MustValidateInput(w, obj)
exists, account, err := GetUserByLogin(obj.Email)
if err != nil {
- panic(err)
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ return
}
if exists && account != nil {
log.Warn("email already exists")
- //obj.ID = account.ID
- panic("email already existed")
+ api.WriteError(w, errEmailAlreadyExists, http.StatusConflict)
+ return
} else {
obj.ID = getUIDByEmail(obj.Email)
}
randStr := util.GenerateSecureString(8)
- hash, err := bcrypt.GenerateFromPassword([]byte(randStr), bcrypt.DefaultCost)
- if err != nil {
- panic(err)
+ if err := security.SetPassword(obj, randStr); err != nil {
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ return
}
- obj.Password = string(hash)
-
ctx := orm.NewContextWithParent(req.Context())
ctx.Refresh = orm.WaitForRefresh
err = orm.Save(ctx, obj)
if err != nil {
- panic(err)
+ api.WriteError(w, err.Error(), http.StatusInternalServerError)
+ return
}
obj.Password = randStr
+ // The one-time bootstrap password should be returned to the caller, but the
+ // persisted verifier material must stay server-side only.
+ obj.PasswordSalt = ""
+ obj.PasswordVerifier = ""
api.WriteJSON(w, obj, 200)
}
+
+func validateSecurePassword(password string) error {
+ if util.ValidateSecure(password) {
+ return nil
+ }
+ return cerr.NewWithHTTPCode(http.StatusBadRequest, errInsecurePassword)
+}
+
+func resolveUserByLogin(login string, items []security.UserAccount) (bool, *security.UserAccount, error) {
+ switch len(items) {
+ case 0:
+ return false, nil, nil
+ case 1:
+ return true, &items[0], nil
+ default:
+ log.Warnf("invalid users, more than one account was associated with the same email: %v", login)
+ return false, nil, fmt.Errorf("multiple accounts found for login %q", login)
+ }
+}
diff --git a/modules/security/native/user_test.go b/modules/security/native/user_test.go
new file mode 100644
index 000000000..f1e8fc82e
--- /dev/null
+++ b/modules/security/native/user_test.go
@@ -0,0 +1,62 @@
+// Copyright (C) INFINI Labs & INFINI LIMITED.
+//
+// The INFINI Framework is offered under the GNU Affero General Public License v3.0
+// and as commercial software.
+//
+// For commercial licensing, contact us at:
+// - Website: infinilabs.com
+// - Email: hello@infini.ltd
+//
+// Open Source licensed under AGPL V3:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package native
+
+import (
+ "strings"
+ "testing"
+
+ "infini.sh/framework/core/security"
+)
+
+// Weak passwords should now fail as normal validation errors instead of aborting
+// the request flow via panic.
+func TestValidateSecurePassword(t *testing.T) {
+ if err := validateSecurePassword("weak"); err == nil {
+ t.Fatal("expected weak password to be rejected")
+ }
+
+ if err := validateSecurePassword("StrongPassw0rd!"); err != nil {
+ t.Fatalf("expected strong password to pass validation, got %v", err)
+ }
+}
+
+func TestResolveUserByLogin(t *testing.T) {
+ found, user, err := resolveUserByLogin("missing@example.org", nil)
+ if err != nil || found || user != nil {
+ t.Fatalf("expected empty result for missing user, got found=%v user=%#v err=%v", found, user, err)
+ }
+
+ items := []security.UserAccount{{}}
+ items[0].Email = "admin@example.org"
+ found, user, err = resolveUserByLogin("admin@example.org", items)
+ if err != nil || !found || user == nil || user.Email != "admin@example.org" {
+ t.Fatalf("expected single user match, got found=%v user=%#v err=%v", found, user, err)
+ }
+
+ _, _, err = resolveUserByLogin("dup@example.org", []security.UserAccount{{}, {}})
+ if err == nil || !strings.Contains(err.Error(), "multiple accounts found") {
+ t.Fatalf("expected duplicate login error, got %v", err)
+ }
+}
diff --git a/plugins/badger/badger.go b/plugins/badger/badger.go
index f18eaaf38..05d4752f5 100644
--- a/plugins/badger/badger.go
+++ b/plugins/badger/badger.go
@@ -283,6 +283,10 @@ func (filter *Module) GetCompressedValue(bucket string, key []byte) ([]byte, err
}
func (filter *Module) AddValueCompress(bucket string, key []byte, value []byte) error {
+ return filter.AddValueCompressWithTTL(bucket, key, value, 0)
+}
+
+func (filter *Module) AddValueCompressWithTTL(bucket string, key []byte, value []byte, ttl time.Duration) error {
value, err := lz4.Encode(nil, value)
if err != nil {
log.Error("Failed to encode:", err)
@@ -291,7 +295,7 @@ func (filter *Module) AddValueCompress(bucket string, key []byte, value []byte)
stats.Increment("badger", bucket+"::add_compress")
- return filter.AddValue(bucket, key, value)
+ return filter.AddValueWithTTL(bucket, key, value, ttl)
}
func joinKey(bucket string, key []byte) []byte {
@@ -299,6 +303,10 @@ func joinKey(bucket string, key []byte) []byte {
}
func (filter *Module) AddValue(bucket string, key []byte, value []byte) error {
+ return filter.AddValueWithTTL(bucket, key, value, 0)
+}
+
+func (filter *Module) AddValueWithTTL(bucket string, key []byte, value []byte, ttl time.Duration) error {
if filter.closed {
return errors.New("module closed")
}
@@ -310,8 +318,10 @@ func (filter *Module) AddValue(bucket string, key []byte, value []byte) error {
}
bkt := filter.getOrInitBucket(bucket)
err := bkt.Update(func(txn *badger.Txn) error {
- err := txn.Set(key, value)
- return err
+ if ttl > 0 {
+ return txn.SetEntry(badger.NewEntry(key, value).WithTTL(ttl))
+ }
+ return txn.Set(key, value)
})
return err
}
diff --git a/plugins/elastic/bulk_indexing/bulk_indexing.go b/plugins/elastic/bulk_indexing/bulk_indexing.go
index f7c1f118a..936fcb9d3 100755
--- a/plugins/elastic/bulk_indexing/bulk_indexing.go
+++ b/plugins/elastic/bulk_indexing/bulk_indexing.go
@@ -78,6 +78,65 @@ type BulkIndexingProcessor struct {
bulkBufferPool *elastic.BulkBufferPool
}
+const bulkLogSampleLimit = 5
+
+func summarizeBulkLogValues(values []string) string {
+ if len(values) == 0 {
+ return "[]"
+ }
+
+ limit := bulkLogSampleLimit
+ if len(values) < limit {
+ limit = len(values)
+ }
+
+ sample := values[:limit]
+ if len(values) > limit {
+ return fmt.Sprintf("%v...(and %d more)", sample, len(values)-limit)
+ }
+
+ return fmt.Sprintf("%v", sample)
+}
+
+func summarizeBulkDetailItem(item elastic.BulkDetailItem) string {
+ parts := make([]string, 0, 2)
+ if len(item.Documents) > 0 {
+ parts = append(parts, fmt.Sprintf("documents=%d sample=%s", len(item.Documents), summarizeBulkLogValues(item.Documents)))
+ }
+ if len(item.Reasons) > 0 {
+ parts = append(parts, fmt.Sprintf("reasons=%d sample=%s", len(item.Reasons), summarizeBulkLogValues(item.Reasons)))
+ }
+ if len(parts) == 0 {
+ return "empty"
+ }
+ return strings.Join(parts, ", ")
+}
+
+func summarizeBulkResult(bulkResult *elastic.BulkResult) string {
+ if bulkResult == nil {
+ return ""
+ }
+
+ parts := []string{
+ fmt.Sprintf(
+ "summary={success:%d invalid:%d failure:%d}",
+ bulkResult.Summary.Success.Count,
+ bulkResult.Summary.Invalid.Count,
+ bulkResult.Summary.Failure.Count,
+ ),
+ fmt.Sprintf("error=%v", bulkResult.Error),
+ fmt.Sprintf("error_msgs=%d sample=%s", len(bulkResult.ErrorMsgs), summarizeBulkLogValues(bulkResult.ErrorMsgs)),
+ fmt.Sprintf("codes=%d", len(bulkResult.Codes)),
+ fmt.Sprintf("indices=%d", len(bulkResult.Indices)),
+ fmt.Sprintf("actions=%d", len(bulkResult.Actions)),
+ fmt.Sprintf("detail={failure:%s, invalid:%s}", summarizeBulkDetailItem(bulkResult.Detail.Failure), summarizeBulkDetailItem(bulkResult.Detail.Invalid)),
+ }
+
+ return strings.Join(parts, ", ")
+}
+
+var queueOwners sync.Map
+
type Config struct {
NumOfSlices int `config:"num_of_slices"`
Slices []int `config:"slices"`
@@ -233,7 +292,21 @@ func (processor *BulkIndexingProcessor) Process(c *pipeline.Context) error {
log.Error("error in bulk indexing processor,", v)
}
}
- log.Debug("exit bulk indexing processor")
+ if processor.bulkStats != nil {
+ logFn := log.Tracef
+ if processor.bulkStats.Summary.Invalid.Count > 0 || processor.bulkStats.Summary.Failure.Count > 0 || len(processor.bulkStats.ErrorMsgs) > 0 {
+ logFn = log.Debugf
+ }
+ logFn(
+ "exit bulk indexing processor, success=%d, invalid=%d, failure=%d, error_msgs=%d",
+ processor.bulkStats.Summary.Success.Count,
+ processor.bulkStats.Summary.Invalid.Count,
+ processor.bulkStats.Summary.Failure.Count,
+ len(processor.bulkStats.ErrorMsgs),
+ )
+ } else {
+ log.Trace("exit bulk indexing processor")
+ }
}()
//handle updates
@@ -264,10 +337,11 @@ func (processor *BulkIndexingProcessor) Process(c *pipeline.Context) error {
}
}
processor.detectorRunning = false
- log.Debug("exit detector for active queue")
+ log.Trace("exit detector for active queue")
processor.wg.Done()
}()
+ lastDispatch := time.Now()
for {
if global.ShuttingDown() {
@@ -299,11 +373,12 @@ func (processor *BulkIndexingProcessor) Process(c *pipeline.Context) error {
}
//if have depth and not in in flight
if !processor.config.SkipEmptyQueue || queue.HasLag(v) {
- _, ok := processor.inFlightQueueConfigs.Load(v.ID)
+ ok := processor.hasInFlightQueue(v.ID)
if !ok {
if global.Env().IsDebug {
log.Tracef("detecting new queue: %v", v.Name)
}
+ lastDispatch = time.Now()
processor.HandleQueueConfig(v, c)
}
} else {
@@ -315,12 +390,35 @@ func (processor *BulkIndexingProcessor) Process(c *pipeline.Context) error {
if processor.config.DetectIntervalInMs > 0 {
time.Sleep(time.Millisecond * time.Duration(processor.config.DetectIntervalInMs))
}
+ if shouldQuitActiveQueueDetection(
+ lastDispatch,
+ time.Duration(processor.config.IdleTimeoutInSecond)*time.Second,
+ time.Duration(processor.config.DetectIntervalInMs)*time.Millisecond,
+ util.MapLength(&processor.inFlightQueueConfigs),
+ ) {
+ if processor.bulkStats != nil {
+ logFn := log.Tracef
+ if processor.bulkStats.Summary.Invalid.Count > 0 || processor.bulkStats.Summary.Failure.Count > 0 || len(processor.bulkStats.ErrorMsgs) > 0 {
+ logFn = log.Debugf
+ }
+ logFn(
+ "active queue detector idle exit, success=%d, invalid=%d, failure=%d, inflight=%d",
+ processor.bulkStats.Summary.Success.Count,
+ processor.bulkStats.Summary.Invalid.Count,
+ processor.bulkStats.Summary.Failure.Count,
+ util.MapLength(&processor.inFlightQueueConfigs),
+ )
+ }
+ return
+ }
}
}(c)
}
} else {
cfgs := queue.GetConfigBySelector(&processor.config.Selector)
- log.Debugf("filter queue by:%v, num of queues:%v", processor.config.Selector.ToString(), len(cfgs))
+ if global.Env().IsDebug {
+ log.Tracef("filter queue by:%v, num of queues:%v", processor.config.Selector.ToString(), len(cfgs))
+ }
for _, v := range cfgs {
if global.Env().IsDebug {
log.Tracef("checking queue: %v", v)
@@ -334,14 +432,33 @@ func (processor *BulkIndexingProcessor) Process(c *pipeline.Context) error {
return nil
}
+func shouldQuitActiveQueueDetection(lastDispatch time.Time, idleDuration time.Duration, detectInterval time.Duration, inflight int) bool {
+ if idleDuration <= 0 {
+ return false
+ }
+ if detectInterval < 0 {
+ detectInterval = 0
+ }
+ return inflight == 0 && time.Since(lastDispatch) >= idleDuration+detectInterval
+}
+
const queueHandleSingleton = "queue_handler_singleton"
func (processor *BulkIndexingProcessor) HandleQueueConfig(v *queue.QueueConfig, parentContext *pipeline.Context) {
+ if !processor.acquireQueueOwner(v.ID) {
+ if rate.GetRateLimiter("bulk_queue_owner", v.ID, 1, 1, 30*time.Second).Allow() {
+ log.Debugf("skip queue:[%v], already owned by another local bulk processor", v.ID)
+ }
+ return
+ }
+ defer processor.releaseQueueOwnerIfIdle(v.ID)
//TODO, add config to enable/disable singleton, may have performance issue
ok, _ := locker.Hold(queueHandleSingleton, v.ID, global.Env().SystemConfig.NodeConfig.ID, 60*time.Second, true)
if !ok {
- log.Debugf("failed to hold lock for queue:[%v], already hold by somewhere", v.ID)
+ if rate.GetRateLimiter("bulk_queue_lock", v.ID, 1, 1, 30*time.Second).Allow() {
+ log.Debugf("failed to hold lock for queue:[%v], already hold by somewhere", v.ID)
+ }
return
}
@@ -433,6 +550,10 @@ func (processor *BulkIndexingProcessor) HandleQueueConfig(v *queue.QueueConfig,
}
func (processor *BulkIndexingProcessor) NewBulkWorker(parentContext *pipeline.Context, qConfig *queue.QueueConfig, preferedHost string) {
+ if global.Env().IsDebug {
+ // current time for monitoring and log
+ log.Debugf("starting bulk worker for queue: %v, host: %v at time: %v", qConfig.Name, preferedHost, time.Now().Format(time.RFC3339))
+ }
bulkSizeInByte := processor.config.BulkConfig.GetBulkSizeInBytes()
//check slice
for sliceID := 0; sliceID < processor.config.NumOfSlices; sliceID++ {
@@ -460,50 +581,110 @@ func (processor *BulkIndexingProcessor) NewBulkWorker(parentContext *pipeline.Co
return
}
- processor.Lock()
- v2, exists := processor.inFlightQueueConfigs.Load(key)
- if exists {
+ var workerID = util.GetUUID()
+ v2, reserved := processor.reserveInFlightQueue(key, workerID)
+ if !reserved {
if global.Env().IsDebug {
log.Tracef("[%v], queue [%v], slice_id:%v has more then one consumer, key:%v,v:%v", preferedHost, qConfig.ID, sliceID, key, v2)
}
- processor.Unlock()
continue
- } else {
- var workerID = util.GetUUID()
- log.Debugf("starting worker:[%v], queue:[%v], slice_id:%v, host:[%v]", workerID, qConfig.Name, sliceID, preferedHost)
-
- ctx1 := &pipeline.Context{}
- ctx1.Set("key", key)
- ctx1.Set("workerID", workerID)
- ctx1.Set("sliceID", sliceID)
- ctx1.Set("numOfSlices", processor.config.NumOfSlices)
- ctx1.Set("tag", preferedHost)
- ctx1.Set("qConfig", qConfig)
- ctx1.Set("host", preferedHost)
- ctx1.Set("bulkSizeInByte", bulkSizeInByte)
- err := processor.pool.Submit(&pipeline.Task{
- Handler: func(ctx *pipeline.Context, v ...interface{}) {
- key := ctx.MustGetString("key")
- workerID := ctx.MustGetString("workerID")
- host := ctx.MustGetString("host")
- sliceID := ctx.MustGetInt("sliceID")
- tag := ctx.MustGetString("tag")
- numOfSlices := ctx.MustGetInt("numOfSlices")
- bulkSizeInByte := ctx.MustGetInt("bulkSizeInByte")
- qConfig := ctx.MustGet("qConfig").(*queue.QueueConfig)
- pCtx := v[0].(*pipeline.Context)
- processor.NewSlicedBulkWorker(pCtx, key, workerID, sliceID, numOfSlices, tag, bulkSizeInByte, qConfig, host)
- },
- Context: ctx1,
- Params: []interface{}{parentContext}, // 也可以在创建任务时设置参数
- })
- processor.Unlock()
- if err != nil {
- panic(err)
- }
- processor.wg.Add(1)
}
+
+ log.Tracef("starting worker:[%v], queue:[%v], slice_id:%v, host:[%v]", workerID, qConfig.Name, sliceID, preferedHost)
+
+ ctx1 := &pipeline.Context{}
+ ctx1.Set("key", key)
+ ctx1.Set("workerID", workerID)
+ ctx1.Set("sliceID", sliceID)
+ ctx1.Set("numOfSlices", processor.config.NumOfSlices)
+ ctx1.Set("tag", preferedHost)
+ ctx1.Set("qConfig", qConfig)
+ ctx1.Set("host", preferedHost)
+ ctx1.Set("bulkSizeInByte", bulkSizeInByte)
+ err := processor.pool.Submit(&pipeline.Task{
+ Handler: func(ctx *pipeline.Context, v ...interface{}) {
+ key := ctx.MustGetString("key")
+ workerID := ctx.MustGetString("workerID")
+ host := ctx.MustGetString("host")
+ sliceID := ctx.MustGetInt("sliceID")
+ tag := ctx.MustGetString("tag")
+ numOfSlices := ctx.MustGetInt("numOfSlices")
+ bulkSizeInByte := ctx.MustGetInt("bulkSizeInByte")
+ qConfig := ctx.MustGet("qConfig").(*queue.QueueConfig)
+ pCtx := v[0].(*pipeline.Context)
+ processor.NewSlicedBulkWorker(pCtx, key, workerID, sliceID, numOfSlices, tag, bulkSizeInByte, qConfig, host)
+ },
+ Context: ctx1,
+ Params: []interface{}{parentContext}, // 也可以在创建任务时设置参数
+ })
+ if err != nil {
+ processor.inFlightQueueConfigs.Delete(key)
+ processor.wg.Done()
+ panic(err)
+ }
+ }
+}
+
+func (processor *BulkIndexingProcessor) reserveInFlightQueue(key, workerID string) (interface{}, bool) {
+ processor.Lock()
+ defer processor.Unlock()
+
+ v, exists := processor.inFlightQueueConfigs.Load(key)
+ if exists {
+ return v, false
}
+
+ processor.inFlightQueueConfigs.Store(key, workerID)
+ processor.wg.Add(1)
+
+ return workerID, true
+}
+
+func (processor *BulkIndexingProcessor) hasInFlightQueue(queueID string) bool {
+ if _, ok := processor.inFlightQueueConfigs.Load(queueID); ok {
+ return true
+ }
+
+ queuePrefix := fmt.Sprintf("%v-", queueID)
+ hasInFlight := false
+ processor.inFlightQueueConfigs.Range(func(key, value interface{}) bool {
+ keyStr, ok := key.(string)
+ if ok && strings.HasPrefix(keyStr, queuePrefix) {
+ hasInFlight = true
+ return false
+ }
+ return true
+ })
+
+ return hasInFlight
+}
+
+func (processor *BulkIndexingProcessor) acquireQueueOwner(queueID string) bool {
+ owner, loaded := queueOwners.LoadOrStore(queueID, processor.id)
+ if !loaded {
+ return true
+ }
+
+ return owner == processor.id
+}
+
+func (processor *BulkIndexingProcessor) releaseQueueOwnerIfIdle(queueID string) {
+ if processor.hasInFlightQueue(queueID) {
+ return
+ }
+
+ owner, ok := queueOwners.Load(queueID)
+ if ok && owner == processor.id {
+ queueOwners.Delete(queueID)
+ }
+}
+
+func isIgnorableAcquireConsumerError(err error) bool {
+ if err == nil {
+ return false
+ }
+
+ return util.ContainStr(err.Error(), "already owning this topic")
}
var xxHashPool = sync.Pool{
@@ -549,8 +730,6 @@ func (processor *BulkIndexingProcessor) getConsumerConfig(queueID, consumerName
}
func (processor *BulkIndexingProcessor) NewSlicedBulkWorker(ctx *pipeline.Context, key, workerID string, sliceID, maxSlices int, tag string, bulkSizeInByte int, qConfig *queue.QueueConfig, host string) {
- processor.inFlightQueueConfigs.Store(key, workerID)
-
defer func() {
if !global.Env().IsDebug {
if r := recover(); r != nil {
@@ -571,6 +750,7 @@ func (processor *BulkIndexingProcessor) NewSlicedBulkWorker(ctx *pipeline.Contex
}
}
processor.inFlightQueueConfigs.Delete(key)
+ processor.releaseQueueOwnerIfIdle(qConfig.ID)
processor.wg.Done()
if global.Env().IsDebug {
log.Tracef("exit slice worker, worker:[%v], queue:%v, slice_id:%v, key:%v", workerID, qConfig.ID, sliceID, key)
@@ -600,12 +780,15 @@ func (processor *BulkIndexingProcessor) NewSlicedBulkWorker(ctx *pipeline.Contex
var consumerInstance queue.ConsumerAPI
consumerInstance, err = queue.AcquireConsumer(qConfig, consumerConfig, workerID)
if err != nil || consumerInstance == nil {
- if util.ContainStr(err.Error(), "already owning this topic") {
+ if isIgnorableAcquireConsumerError(err) {
if global.Env().IsDebug {
- log.Warnf("other consumer already owning this topic, queue:%v-%v, slice_id:%v", qConfig.Name, qConfig.ID, sliceID)
+ log.Warnf("skip duplicate consumer acquisition, queue:%v-%v, slice_id:%v, err:%v", qConfig.Name, qConfig.ID, sliceID, err)
}
return
}
+ if err == nil {
+ err = errors.New("failed to acquire queue consumer")
+ }
panic(err)
}
@@ -681,7 +864,9 @@ func (processor *BulkIndexingProcessor) NewSlicedBulkWorker(ctx *pipeline.Contex
log.Errorf("should not submit this bulk request, worker[%v], queue:[%v], slice:[%v], offset:[%v]->[%v],%v, msg:%v", workerID, qConfig.ID, sliceID, committedOffset, offset, err, mainBuf.GetMessageCount())
}
}
- log.Debugf("exit worker[%v], message count[%d], queue:[%v], slice_id:%v", workerID, mainBuf.GetMessageCount(), qConfig.ID, sliceID)
+ if global.Env().IsDebug {
+ log.Tracef("exit worker[%v], message count[%d], queue:[%v], slice_id:%v", workerID, mainBuf.GetMessageCount(), qConfig.ID, sliceID)
+ }
}()
if global.Env().IsDebug {
@@ -798,7 +983,11 @@ READ_DOCS:
consumerConfig.KeepActive()
messages, timeout, err := consumerInstance.FetchMessages(ctx1, consumerConfig.FetchMaxMessages)
stats.IncrementBy("queue", qConfig.ID+".msg_fetched_from_queue", int64(len(messages)))
- log.Debugf("slice worker, worker:[%v], [%v][%v][%v][%v] fetched message:%v,ctx:%v,timeout:%v,err:%v", workerID, qConfig.Name, consumerConfig.Group, consumerConfig.Name, sliceID, len(messages), ctx1.String(), timeout, err)
+ if err != nil {
+ log.Debugf("slice worker, worker:[%v], [%v][%v][%v][%v] fetched message:%v,ctx:%v,timeout:%v,err:%v", workerID, qConfig.Name, consumerConfig.Group, consumerConfig.Name, sliceID, len(messages), ctx1.String(), timeout, err)
+ } else if len(messages) > 0 {
+ log.Tracef("slice worker, worker:[%v], [%v][%v][%v][%v] fetched message:%v,ctx:%v,timeout:%v,err:%v", workerID, qConfig.Name, consumerConfig.Group, consumerConfig.Name, sliceID, len(messages), ctx1.String(), timeout, err)
+ }
if err != nil {
if strings.Contains(err.Error(), "dirty_read") || err.Error() == "EOF" || err.Error() == "unexpected EOF" {
ctx.CancelTask()
@@ -905,6 +1094,10 @@ READ_DOCS:
mainBuf.WriteByteBuffer(pop.Data)
}
+ // Keep the in-memory offset aligned with the data already buffered.
+ // If the current message triggers an immediate flush, its NextOffset must be committed too.
+ offset = advanceBufferedOffset(pop.NextOffset)
+
if global.Env().IsDebug {
log.Tracef("slice worker, worker:[%v], message count: %v, size: %v", workerID, mainBuf.GetMessageCount(), util.ByteSize(uint64(mainBuf.GetMessageSize())))
}
@@ -949,7 +1142,7 @@ READ_DOCS:
if offset != nil && committedOffset != nil && !offset.Equals(*committedOffset) {
err := consumerInstance.CommitOffset(*offset)
if err != nil {
- log.Errorf("🔧 offset commit failed, worker:[%v], queue:[%v], slice:[%v], offset:[%v], err:%v", workerID, qConfig.Name, sliceID, *offset, err)
+ log.Errorf("offset commit failed, worker:[%v], queue:[%v], slice:[%v], offset:[%v], err:%v", workerID, qConfig.Name, sliceID, *offset, err)
panic(err)
}
@@ -958,27 +1151,18 @@ READ_DOCS:
}
// fix: update committedOffset immediately after successful commit, to ensure state consistency
committedOffset = offset
- log.Debugf("🔧 offset committed successfully, worker:[%v], queue:[%v], slice:[%v], offset:[%v]", workerID, qConfig.Name, sliceID, *offset)
- } else {
if global.Env().IsDebug {
- log.Debugf("🔧 offset not changed, skip commit, worker:[%v], queue:[%v], slice:[%v], offset:[%v], committed:[%v]", workerID, qConfig.Name, sliceID, offset, committedOffset)
+ log.Tracef("offset committed, worker:[%v], queue:[%v], slice:[%v], offset:[%v]", workerID, qConfig.Name, sliceID, *offset)
}
+ } else {
+ // skip unchanged offset silently to avoid noisy debug logs
}
- // fix: this code is moved to loop outside (line 970) to avoid updating offset in the middle of bulk submission
- // offset = &pop.NextOffset
}
} else {
log.Errorf("should not submit this bulk request, worker[%v], queue:[%v], slice:[%v], offset:[%v]->[%v],%v, msg:%v", workerID, qConfig.ID, sliceID, committedOffset, offset, err, msgCount)
}
}
-
- // fix: update offset after each message is processed, to ensure progress sync with actual processing
- // so even if it crashes before submission, it will not repeat processing messages written to the buffer after restart
- offset = &pop.NextOffset
}
-
- // fix: remove this code to avoid overwriting the updated offset in the loop
- // offset = &ctx1.NextOffset
}
if time.Since(lastCommit) > idleDuration && mainBuf.GetMessageSize() > 0 {
@@ -1002,7 +1186,7 @@ CLEAN_BUFFER:
}
if global.Env().IsDebug {
- log.Debugf("cleanup buffer, queue:[%v], slice_id:%v, offset [%v]-[%v], bulk failed (host: %v, err: %v)", qConfig.ID, sliceID, committedOffset, offset, host, err)
+ log.Tracef("cleanup buffer, queue:[%v], slice_id:%v, offset [%v]-[%v], bulk failed (host: %v, err: %v)", qConfig.ID, sliceID, committedOffset, offset, host, err)
}
lastCommit = time.Now()
// check bulk result, if ok, then commit offset, or retry non-200 requests, or save failure offset
@@ -1112,11 +1296,18 @@ func (processor *BulkIndexingProcessor) submitBulkRequest(ctx *pipeline.Context,
if bulkResult != nil {
msg = bulkResult.Detail
}
- log.Warnf("elasticsearch [%v], stats:%v, detail: %v, err:%v", meta.Config.Name, statsMap, msg, err)
+ log.Warnf(
+ "elasticsearch [%v], stats:%v, detail:{failure:%s, invalid:%s}, err:%v",
+ meta.Config.Name,
+ statsMap,
+ summarizeBulkDetailItem(msg.Failure),
+ summarizeBulkDetailItem(msg.Invalid),
+ err,
+ )
}
if global.Env().IsDebug {
- log.Debug(tag, ", ", meta.Config.Name, ", ", host, ", stats: ", statsMap, ", count: ", count, ", size: ", util.ByteSize(uint64(size)), ", elapsed: ", time.Since(start), ", continue: ", continueRequest, ", bulkResult: ", bulkResult)
+ log.Debug(tag, ", ", meta.Config.Name, ", ", host, ", stats: ", statsMap, ", count: ", count, ", size: ", util.ByteSize(uint64(size)), ", elapsed: ", time.Since(start), ", continue: ", continueRequest, ", bulkResult: ", summarizeBulkResult(bulkResult))
} else {
if processor.config.VerboseBulkResult {
log.Info("queue:", qConfig.Name, ", ", meta.Config.Name, ", ", host, ", stats: ", statsMap, ", count: ", count, ", size: ", util.ByteSize(uint64(size)), ", elapsed: ", time.Since(start), ", continue: ", continueRequest)
@@ -1167,6 +1358,11 @@ func appendStrArr(arr []string, size int, elems []string) []string {
return append(arr, elems...)
}
+func advanceBufferedOffset(nextOffset queue.Offset) *queue.Offset {
+ next := nextOffset
+ return &next
+}
+
func (processor *BulkIndexingProcessor) getElasticsearchMetadata(qConfig *queue.QueueConfig) (string, *elastic.ElasticsearchMetadata) {
elasticsearch, ok := qConfig.Labels["elasticsearch"]
diff --git a/plugins/elastic/bulk_indexing/bulk_indexing_test.go b/plugins/elastic/bulk_indexing/bulk_indexing_test.go
index e37b0ffe0..f462d8f34 100644
--- a/plugins/elastic/bulk_indexing/bulk_indexing_test.go
+++ b/plugins/elastic/bulk_indexing/bulk_indexing_test.go
@@ -28,9 +28,13 @@
package bulk_indexing
import (
+ stdErrors "errors"
"github.com/OneOfOne/xxhash"
"github.com/stretchr/testify/assert"
+ "infini.sh/framework/core/queue"
+ "sync"
"testing"
+ "time"
)
func TestXXHash(t *testing.T) {
@@ -84,3 +88,90 @@ func TestXXHash(t *testing.T) {
}
}
+
+func TestReserveInFlightQueue(t *testing.T) {
+ processor := &BulkIndexingProcessor{}
+
+ current, reserved := processor.reserveInFlightQueue("queue-0", "worker-1")
+ assert.True(t, reserved)
+ assert.Equal(t, "worker-1", current)
+
+ stored, exists := processor.inFlightQueueConfigs.Load("queue-0")
+ assert.True(t, exists)
+ assert.Equal(t, "worker-1", stored)
+
+ current, reserved = processor.reserveInFlightQueue("queue-0", "worker-2")
+ assert.False(t, reserved)
+ assert.Equal(t, "worker-1", current)
+
+ processor.inFlightQueueConfigs.Delete("queue-0")
+ processor.wg.Done()
+}
+
+func TestHasInFlightQueue(t *testing.T) {
+ processor := &BulkIndexingProcessor{}
+
+ assert.False(t, processor.hasInFlightQueue("queue-0"))
+
+ processor.inFlightQueueConfigs.Store("queue-0-0", "worker-1")
+ assert.True(t, processor.hasInFlightQueue("queue-0"))
+
+ processor.inFlightQueueConfigs.Delete("queue-0-0")
+ assert.False(t, processor.hasInFlightQueue("queue-0"))
+}
+
+func TestAcquireQueueOwner(t *testing.T) {
+ queueOwners = sync.Map{}
+
+ processor1 := &BulkIndexingProcessor{id: "processor-1"}
+ processor2 := &BulkIndexingProcessor{id: "processor-2"}
+
+ assert.True(t, processor1.acquireQueueOwner("queue-0"))
+ assert.True(t, processor1.acquireQueueOwner("queue-0"))
+ assert.False(t, processor2.acquireQueueOwner("queue-0"))
+
+ queueOwners = sync.Map{}
+}
+
+func TestReleaseQueueOwnerIfIdle(t *testing.T) {
+ queueOwners = sync.Map{}
+
+ processor := &BulkIndexingProcessor{id: "processor-1"}
+ assert.True(t, processor.acquireQueueOwner("queue-0"))
+
+ processor.inFlightQueueConfigs.Store("queue-0-0", "worker-1")
+ processor.releaseQueueOwnerIfIdle("queue-0")
+ _, exists := queueOwners.Load("queue-0")
+ assert.True(t, exists)
+
+ processor.inFlightQueueConfigs.Delete("queue-0-0")
+ processor.releaseQueueOwnerIfIdle("queue-0")
+ _, exists = queueOwners.Load("queue-0")
+ assert.False(t, exists)
+}
+
+func TestIsIgnorableAcquireConsumerError(t *testing.T) {
+ assert.True(t, isIgnorableAcquireConsumerError(stdErrors.New("already owning this topic")))
+ assert.False(t, isIgnorableAcquireConsumerError(stdErrors.New("the consumer is in fighting list")))
+ assert.False(t, isIgnorableAcquireConsumerError(stdErrors.New("some other error")))
+ assert.False(t, isIgnorableAcquireConsumerError(nil))
+}
+
+func TestShouldQuitActiveQueueDetection(t *testing.T) {
+ assert.False(t, shouldQuitActiveQueueDetection(time.Now(), 5*time.Second, 5*time.Second, 0))
+ assert.False(t, shouldQuitActiveQueueDetection(time.Now().Add(-10*time.Second), 5*time.Second, 5*time.Second, 1))
+ assert.False(t, shouldQuitActiveQueueDetection(time.Now().Add(-9*time.Second), 5*time.Second, 5*time.Second, 0))
+ assert.True(t, shouldQuitActiveQueueDetection(time.Now().Add(-10*time.Second), 5*time.Second, 5*time.Second, 0))
+ assert.True(t, shouldQuitActiveQueueDetection(time.Now().Add(-5*time.Second), 5*time.Second, 0, 0))
+}
+
+func TestAdvanceBufferedOffsetUsesCurrentMessageNextOffset(t *testing.T) {
+ previousCommitted := queue.NewOffsetWithVersion(0, 100, 1)
+ currentNext := queue.NewOffsetWithVersion(0, 200, 1)
+
+ offset := advanceBufferedOffset(currentNext)
+
+ assert.NotNil(t, offset)
+ assert.True(t, offset.Equals(currentNext))
+ assert.False(t, offset.Equals(previousCommitted))
+}
diff --git a/plugins/queue/consumer/consumer.go b/plugins/queue/consumer/consumer.go
index 5d768d19c..9f9d2bc8f 100755
--- a/plugins/queue/consumer/consumer.go
+++ b/plugins/queue/consumer/consumer.go
@@ -28,6 +28,7 @@ import (
"infini.sh/framework/core/errors"
"infini.sh/framework/core/locker"
"runtime"
+ "strings"
"sync"
"time"
@@ -191,23 +192,48 @@ func (processor *QueueConsumerProcessor) Name() string {
return name
}
+func getRecoveredMessage(r interface{}) string {
+ switch v := r.(type) {
+ case error:
+ return v.Error()
+ case runtime.Error:
+ return v.Error()
+ case string:
+ return v
+ default:
+ return fmt.Sprint(v)
+ }
+}
+
+func isExpectedQueueShutdownPanic(message string, contexts ...*pipeline.Context) bool {
+ normalized := strings.ToLower(strings.TrimSpace(message))
+ if normalized == "" || !strings.Contains(normalized, "module closed") {
+ return false
+ }
+ if global.ShuttingDown() {
+ return true
+ }
+ for _, ctx := range contexts {
+ if ctx != nil && (ctx.IsCanceled() || ctx.IsFailed()) {
+ return true
+ }
+ }
+ return false
+}
+
func (processor *QueueConsumerProcessor) Process(c *pipeline.Context) error {
defer func() {
if !global.Env().IsDebug {
if r := recover(); r != nil {
- var v string
- switch r.(type) {
- case error:
- v = r.(error).Error()
- case runtime.Error:
- v = r.(runtime.Error).Error()
- case string:
- v = r.(string)
+ v := getRecoveredMessage(r)
+ if isExpectedQueueShutdownPanic(v, c) {
+ log.Debug("queue consumer processor stopped during shutdown,", v)
+ } else {
+ log.Error("error in consumer processor,", v)
}
- log.Error("error in consumer processor,", v)
}
}
- log.Debug("exit consumer processor")
+ log.Trace("exit consumer processor")
}()
//handle updates
@@ -221,20 +247,16 @@ func (processor *QueueConsumerProcessor) Process(c *pipeline.Context) error {
defer func() {
if !global.Env().IsDebug {
if r := recover(); r != nil {
- var v string
- switch r.(type) {
- case error:
- v = r.(error).Error()
- case runtime.Error:
- v = r.(runtime.Error).Error()
- case string:
- v = r.(string)
+ v := getRecoveredMessage(r)
+ if isExpectedQueueShutdownPanic(v, c) {
+ log.Debug("queue processor stopped during shutdown,", v)
+ } else {
+ log.Error("error in queue processor,", v)
}
- log.Error("error in queue processor,", v)
}
}
processor.detectorRunning = false
- log.Debug("exit detector for active queue")
+ log.Trace("exit detector for active queue")
processor.wg.Done()
}()
@@ -285,7 +307,7 @@ func (processor *QueueConsumerProcessor) Process(c *pipeline.Context) error {
log.Tracef("quite detect after idle for %v ms", processor.config.QuitDetectAfterIdleInMs)
inflight := util.MapLength(&processor.inFlightQueueConfigs)
if inflight == 0 {
- log.Debugf("quite detect after idle for %v ms, inflight: %v", processor.config.QuitDetectAfterIdleInMs, inflight)
+ log.Tracef("quite detect after idle for %v ms, inflight: %v", processor.config.QuitDetectAfterIdleInMs, inflight)
return
}
}
@@ -358,7 +380,7 @@ func (processor *QueueConsumerProcessor) HandleQueueConfig(qConfig *queue.QueueC
continue
} else {
var workerID = util.GetUUID()
- log.Debugf("starting worker:[%v], queue:[%v], slice_id:%v", workerID, qConfig.Name, sliceID)
+ log.Tracef("starting worker:[%v], queue:[%v], slice_id:%v", workerID, qConfig.Name, sliceID)
processor.wg.Add(1)
contextForWorker := pipeline.Context{}
@@ -412,16 +434,12 @@ func (processor *QueueConsumerProcessor) NewSlicedWorker(ctx *pipeline.Context,
defer func() {
if !global.Env().IsDebug {
if r := recover(); r != nil {
- var v string
- switch r.(type) {
- case error:
- v = r.(error).Error()
- case runtime.Error:
- v = r.(runtime.Error).Error()
- case string:
- v = r.(string)
+ v := getRecoveredMessage(r)
+ if isExpectedQueueShutdownPanic(v, ctx, parentContext) {
+ log.Debugf("consumer processor stopped during shutdown, queue:%v, slice_id:%v, %v", qConfig.ID, sliceID, v)
+ } else {
+ log.Errorf("error in consumer processor, %v, queue:%v, slice_id:%v", v, qConfig.ID, sliceID)
}
- log.Errorf("error in consumer processor, %v, queue:%v, slice_id:%v", v, qConfig.ID, sliceID)
}
}
processor.inFlightQueueConfigs.Delete(key)
@@ -473,24 +491,18 @@ func (processor *QueueConsumerProcessor) NewSlicedWorker(ctx *pipeline.Context,
defer xxHashPool.Put(xxHash)
defer func() {
- defer log.Debugf("exit worker[%v], queue:[%v], slice_id:%v", workerID, qConfig.ID, sliceID)
+ defer log.Tracef("exit worker[%v], queue:[%v], slice_id:%v", workerID, qConfig.ID, sliceID)
if !global.Env().IsDebug {
if r := recover(); r != nil {
- var v string
- switch r.(type) {
- case error:
- v = r.(error).Error()
- case runtime.Error:
- v = r.(runtime.Error).Error()
- case string:
- v = r.(string)
- }
- if v != "empty queue" {
+ v := getRecoveredMessage(r)
+ if isExpectedQueueShutdownPanic(v, ctx, parentContext) {
+ log.Debugf("worker[%v], queue:[%v], slice:[%v] stopped during shutdown, offset:[%v]->[%v], %v", workerID, qConfig.ID, sliceID, initOffset, offset, v)
+ } else if v != "empty queue" {
log.Errorf("worker[%v], queue:[%v], slice:[%v], offset:[%v]->[%v],%v", workerID, qConfig.ID, sliceID, initOffset, offset, v)
ctx.Failed(fmt.Errorf("panic in slice worker: %+v", r))
- }
- if parentContext != nil {
- parentContext.RecordError(fmt.Errorf("panic in slice worker: %+v", r))
+ if parentContext != nil {
+ parentContext.RecordError(fmt.Errorf("panic in slice worker: %+v", r))
+ }
}
}
}
@@ -585,8 +597,10 @@ READ_DOCS:
}
consumerConfig.KeepActive()
messages, timeout, err := consumerInstance.FetchMessages(ctx1, consumerConfig.FetchMaxMessages)
- if global.Env().IsDebug {
+ if err != nil {
log.Debugf("[%v] slice_worker, [%v][%v] consume message:%v,ctx:%v,timeout:%v,err:%v", qConfig.Name, consumerConfig.Name, sliceID, len(messages), ctx1.String(), timeout, err)
+ } else if global.Env().IsDebug && len(messages) > 0 {
+ log.Tracef("[%v] slice_worker, [%v][%v] consume message:%v,ctx:%v,timeout:%v,err:%v", qConfig.Name, consumerConfig.Name, sliceID, len(messages), ctx1.String(), timeout, err)
}
if err != nil {
@@ -707,12 +721,12 @@ CLEAN_BUFFER:
if processor.config.QuitNeedTag && processor.config.QuitNeedTagName != "" && !ctx.HasTag(processor.config.QuitNeedTagName) {
time.Sleep(1 * time.Second)
- log.Debug("EOF without quit tag, sleep 1s: ", qConfig.Name)
+ log.Trace("EOF without quit tag, sleep 1s: ", qConfig.Name)
goto READ_DOCS
}
ctx.CancelTask()
- log.Debug("EOF, cancel task: ", qConfig.Name)
+ log.Trace("EOF, cancel task: ", qConfig.Name)
return
}
diff --git a/plugins/queue/consumer/consumer_test.go b/plugins/queue/consumer/consumer_test.go
new file mode 100644
index 000000000..b6fbda166
--- /dev/null
+++ b/plugins/queue/consumer/consumer_test.go
@@ -0,0 +1,37 @@
+package consumer
+
+import (
+ "context"
+ "testing"
+
+ "infini.sh/framework/core/pipeline"
+)
+
+func TestIsExpectedQueueShutdownPanicRequiresShutdownSignal(t *testing.T) {
+ ctx := &pipeline.Context{Context: context.Background()}
+ if isExpectedQueueShutdownPanic("module closed", ctx) {
+ t.Fatal("expected module closed without shutdown or cancellation to remain an error")
+ }
+ if isExpectedQueueShutdownPanic("boom", ctx) {
+ t.Fatal("expected unrelated panic message to remain an error")
+ }
+}
+
+func TestIsExpectedQueueShutdownPanicTreatsCanceledContextAsExpected(t *testing.T) {
+ baseCtx, cancel := context.WithCancel(context.Background())
+ ctx := &pipeline.Context{Context: baseCtx}
+ cancel()
+
+ if !isExpectedQueueShutdownPanic("module closed", ctx) {
+ t.Fatal("expected module closed during context cancellation to be treated as shutdown noise")
+ }
+}
+
+func TestGetRecoveredMessage(t *testing.T) {
+ if got := getRecoveredMessage("boom"); got != "boom" {
+ t.Fatalf("unexpected string recovery message: %q", got)
+ }
+ if got := getRecoveredMessage(context.Canceled); got != context.Canceled.Error() {
+ t.Fatalf("unexpected error recovery message: %q", got)
+ }
+}
diff --git a/plugins/simple_kv/simple.go b/plugins/simple_kv/simple.go
index f5510bf75..d5e7aad1b 100644
--- a/plugins/simple_kv/simple.go
+++ b/plugins/simple_kv/simple.go
@@ -30,6 +30,7 @@ package simple_kv
import (
"errors"
"sync"
+ "time"
"github.com/bkaradzic/go-lz4"
log "github.com/cihub/seelog"
@@ -108,6 +109,11 @@ func (filter *SimpleKV) GetCompressedValue(bucket string, key []byte) ([]byte, e
}
func (filter *SimpleKV) AddValueCompress(bucket string, key []byte, value []byte) error {
+ return filter.AddValueCompressWithTTL(bucket, key, value, 0)
+}
+
+func (filter *SimpleKV) AddValueCompressWithTTL(bucket string, key []byte, value []byte, ttl time.Duration) error {
+ _ = ttl
value, err := lz4.Encode(nil, value)
if err != nil {
log.Error("Failed to encode:", err)
@@ -122,6 +128,11 @@ func joinKey(bucket string, key []byte) string {
}
func (filter *SimpleKV) AddValue(bucket string, key []byte, value []byte) error {
+ return filter.AddValueWithTTL(bucket, key, value, 0)
+}
+
+func (filter *SimpleKV) AddValueWithTTL(bucket string, key []byte, value []byte, ttl time.Duration) error {
+ _ = ttl
if filter.closed {
return errors.New("module closed")
}