diff --git a/cmd/server/main.go b/cmd/server/main.go index d30d13b..bf2472e 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -2,16 +2,13 @@ package main import ( "context" - "fmt" "log/slog" "os" "github.com/WangYihang/Platypus/pkg/config" "github.com/WangYihang/Platypus/pkg/dependencies" - "github.com/WangYihang/Platypus/pkg/listeners" "github.com/WangYihang/Platypus/pkg/models" "github.com/WangYihang/Platypus/pkg/options" - "github.com/google/uuid" "go.uber.org/fx" "go.uber.org/zap" ) @@ -42,20 +39,17 @@ func main() { lc.Append(fx.Hook{ OnStart: func(context.Context) error { logger.Info("config loaded", zap.Any("config", cfg)) - for _, l := range cfg.Listeners { - switch l.Type { - case listeners.ListenerTypePlainShell: - go listeners.NewPlainShellListener(l.BindHost, l.BindPort).Start(logger) - case listeners.ListenerTypeEncryptedShell: - go listeners.NewEncryptedShellListener(l.BindHost, l.BindPort).Start(logger) - case listeners.ListenerTypeRESTful: - token := uuid.New().String() - logger.Info("using generated uuid as token", zap.String("token", token)) - go listeners.NewRESTfulListener(l.BindHost, l.BindPort, token).Start(logger) - default: - logger.Error("unsupported listener type", zap.String("type", string(l.Type))) - return fmt.Errorf("unsupported listener type: %s", l.Type) - } + logger.Info("starting plain listeners") + for _, listener := range cfg.PlainListeners { + go listener.Start(logger) + } + logger.Info("starting encrypted listeners") + for _, listener := range cfg.EncryptedListeners { + go listener.Start(logger) + } + logger.Info("starting RESTful listeners") + for _, listener := range cfg.RestfulListeners { + go listener.Start(logger) } return nil }, diff --git a/go.mod b/go.mod index 570c3d5..5250c62 100644 --- a/go.mod +++ b/go.mod @@ -9,6 +9,7 @@ require ( github.com/google/uuid v1.3.0 github.com/hashicorp/golang-lru/v2 v2.0.7 github.com/jessevdk/go-flags v1.6.1 + github.com/matishsiao/goInfo v0.0.0-20240924010139-10388a85396f github.com/sevlyar/go-daemon v0.1.5 github.com/shirou/gopsutil v3.21.11+incompatible go.uber.org/fx v1.22.2 diff --git a/go.sum b/go.sum index a2fb53e..e93a195 100644 --- a/go.sum +++ b/go.sum @@ -62,6 +62,8 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI= +github.com/matishsiao/goInfo v0.0.0-20240924010139-10388a85396f h1:XDrsC/9hdgiU9ecceSmYsS2E3fBtFiYc34dAMFgegnM= +github.com/matishsiao/goInfo v0.0.0-20240924010139-10388a85396f/go.mod h1:aEt7p9Rvh67BYApmZwNDPpgircTO2kgdmDUoF/1QmwA= 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/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= diff --git a/pkg/config/config.go b/pkg/config/config.go index c07d449..a817abe 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -13,7 +13,9 @@ import ( // Config represents the configuration. type Config struct { - Listeners []listeners.Listener `json:"listeners" yaml:"listeners" toml:"listeners"` + PlainListeners []listeners.PlainListener `json:"plain" yaml:"plain" toml:"plain"` + EncryptedListeners []listeners.EncryptedListener `json:"encrypted" yaml:"encrypted" toml:"encrypted"` + RestfulListeners []listeners.RESTfulListener `json:"restful" yaml:"restful" toml:"restful"` } // LoadConfig loads the configuration from the given path. @@ -24,7 +26,7 @@ func LoadConfig(path string) (*Config, error) { return loadJSONConfig(path) case ".yaml", ".yml": return loadYAMLConfig(path) - case ".toml": + case ".toml", ".tml": return loadTOMLConfig(path) default: return nil, fmt.Errorf("unsupported config file format: %s", ext) diff --git a/pkg/controllers/status.go b/pkg/controllers/status.go index a9d32b6..95eb50e 100644 --- a/pkg/controllers/status.go +++ b/pkg/controllers/status.go @@ -11,7 +11,7 @@ import ( // NewStatusController returns a new status controller // Status controller returns the status of the server, including hostname, CPU usage, disk usage, and memory usage func NewStatusController() func(c *gin.Context) { - cache := expirable.NewLRU[string, models.Status](1, nil, time.Second*60) + cache := expirable.NewLRU[string, models.Status](1, nil, time.Minute) return func(c *gin.Context) { r, ok := cache.Get("status") if ok { diff --git a/pkg/listeners/encrypted.go b/pkg/listeners/encrypted.go new file mode 100644 index 0000000..ad4f637 --- /dev/null +++ b/pkg/listeners/encrypted.go @@ -0,0 +1,23 @@ +package listeners + +import "go.uber.org/zap" + +// EncryptedListener represents an encrypted shell listener. +type EncryptedListener struct { + commonListener +} + +// NewEncryptedListener creates a new encrypted shell listener. +func NewEncryptedListener(host string, port uint16) *EncryptedListener { + return &EncryptedListener{ + commonListener: commonListener{ + BindHost: host, + BindPort: port, + }, + } +} + +// Start starts the encrypted shell listener. +func (l *EncryptedListener) Start(logger *zap.Logger) { + logger.Info("starting encrypted listener", zap.String("host", l.BindHost), zap.Uint16("port", l.BindPort)) +} diff --git a/pkg/listeners/encrypted_shell_listener.go b/pkg/listeners/encrypted_shell_listener.go deleted file mode 100644 index 1bd5bd3..0000000 --- a/pkg/listeners/encrypted_shell_listener.go +++ /dev/null @@ -1,23 +0,0 @@ -package listeners - -import "go.uber.org/zap" - -// EncryptedShellListener represents an encrypted shell listener. -type EncryptedShellListener struct { - Listener -} - -// NewEncryptedShellListener creates a new encrypted shell listener. -func NewEncryptedShellListener(host string, port uint16) *EncryptedShellListener { - return &EncryptedShellListener{ - Listener: Listener{ - BindHost: host, - BindPort: port, - }, - } -} - -// Start starts the encrypted shell listener. -func (l *EncryptedShellListener) Start(logger *zap.Logger) { - logger.Info("starting encrypted listener", zap.String("host", l.BindHost), zap.Uint16("port", l.BindPort)) -} diff --git a/pkg/listeners/listener.go b/pkg/listeners/listener.go index 6a9e016..7c8178c 100644 --- a/pkg/listeners/listener.go +++ b/pkg/listeners/listener.go @@ -1,20 +1,6 @@ package listeners -// ListenerType represents the type of listener. -type ListenerType string - -const ( - // ListenerTypePlainShell represents a plain shell listener. - ListenerTypePlainShell ListenerType = "plain_shell" - // ListenerTypeEncryptedShell represents an encrypted shell listener. - ListenerTypeEncryptedShell ListenerType = "encrypted_shell" - // ListenerTypeRESTful represents a RESTful listener. - ListenerTypeRESTful ListenerType = "restful" -) - -// Listener is a struct that represents a listener -type Listener struct { - BindHost string `json:"bind_host" yaml:"bind_host" toml:"bind_host"` - BindPort uint16 `json:"bind_port" yaml:"bind_port" toml:"bind_port"` - Type ListenerType `json:"type" yaml:"type" toml:"type"` +type commonListener struct { + BindHost string `json:"bind_host" yaml:"bind_host" toml:"bind_host"` + BindPort uint16 `json:"bind_port" yaml:"bind_port" toml:"bind_port"` } diff --git a/pkg/listeners/plain.go b/pkg/listeners/plain.go new file mode 100644 index 0000000..3c471d4 --- /dev/null +++ b/pkg/listeners/plain.go @@ -0,0 +1,23 @@ +package listeners + +import "go.uber.org/zap" + +// PlainListener represents a plain shell listener. +type PlainListener struct { + commonListener +} + +// NewPlainListener creates a new plain shell listener. +func NewPlainListener(host string, port uint16) *PlainListener { + return &PlainListener{ + commonListener: commonListener{ + BindHost: host, + BindPort: port, + }, + } +} + +// Start starts the plain shell listener. +func (l *PlainListener) Start(logger *zap.Logger) { + logger.Info("starting plain listener", zap.String("host", l.BindHost), zap.Uint16("port", l.BindPort)) +} diff --git a/pkg/listeners/plain_shell_listener.go b/pkg/listeners/plain_shell_listener.go deleted file mode 100644 index 8741cd5..0000000 --- a/pkg/listeners/plain_shell_listener.go +++ /dev/null @@ -1,23 +0,0 @@ -package listeners - -import "go.uber.org/zap" - -// PlainShellListener represents a plain shell listener. -type PlainShellListener struct { - Listener -} - -// NewPlainShellListener creates a new plain shell listener. -func NewPlainShellListener(host string, port uint16) *PlainShellListener { - return &PlainShellListener{ - Listener: Listener{ - BindHost: host, - BindPort: port, - }, - } -} - -// Start starts the plain shell listener. -func (l *PlainShellListener) Start(logger *zap.Logger) { - logger.Info("starting plain listener", zap.String("host", l.BindHost), zap.Uint16("port", l.BindPort)) -} diff --git a/pkg/listeners/restful_listener.go b/pkg/listeners/restful.go similarity index 88% rename from pkg/listeners/restful_listener.go rename to pkg/listeners/restful.go index 29e606d..14a348a 100644 --- a/pkg/listeners/restful_listener.go +++ b/pkg/listeners/restful.go @@ -10,14 +10,14 @@ import ( // RESTfulListener represents a RESTful listener. type RESTfulListener struct { - Listener + commonListener Token string `json:"token" yaml:"token" toml:"token"` } // NewRESTfulListener creates a new RESTful listener. func NewRESTfulListener(host string, port uint16, token string) *RESTfulListener { return &RESTfulListener{ - Listener: Listener{ + commonListener: commonListener{ BindHost: host, BindPort: port, }, @@ -30,6 +30,7 @@ func (l *RESTfulListener) Start(logger *zap.Logger) { logger.Info("starting RESTful listener", zap.String("host", l.BindHost), zap.Uint16("port", l.BindPort)) gin.SetMode(gin.ReleaseMode) r := gin.Default() + logger.Info("configuring routes with token", zap.String("token", l.Token)) routes.ConfigureRoutes(r, logger, l.Token) err := r.Run(fmt.Sprintf("%s:%d", l.BindHost, l.BindPort)) if err != nil { diff --git a/pkg/middlewares/auth.go b/pkg/middlewares/auth.go index 1ac3e31..b19927f 100644 --- a/pkg/middlewares/auth.go +++ b/pkg/middlewares/auth.go @@ -1,16 +1,37 @@ package middlewares import ( + "strings" + "github.com/gin-gonic/gin" ) -// AuthMiddleware is a middleware that checks if the request has the correct token. -func AuthMiddleware(token string) gin.HandlerFunc { +// AuthMiddleware is a middleware that checks if the request has the correct Bearer token. +func AuthMiddleware(expectedToken string) gin.HandlerFunc { return func(c *gin.Context) { - if c.GetHeader("Authorization") != token { + // Get the Authorization header + authorizationHeader := c.GetHeader("Authorization") + if authorizationHeader == "" { + c.AbortWithStatusJSON(401, gin.H{"error": "missing authorization header"}) + return + } + + // Check if the header starts with "Bearer " and extract the token + if !strings.HasPrefix(authorizationHeader, "Bearer ") { + c.AbortWithStatusJSON(401, gin.H{"error": "invalid authorization scheme"}) + return + } + + // Extract the token + token := strings.TrimPrefix(authorizationHeader, "Bearer ") + + // Check if the token matches the expected token + if token != expectedToken { c.AbortWithStatusJSON(401, gin.H{"error": "unauthorized"}) return } + + // Token is valid; proceed to the next handler c.Next() } } diff --git a/pkg/models/status.go b/pkg/models/status.go index 4388167..e15cdf7 100644 --- a/pkg/models/status.go +++ b/pkg/models/status.go @@ -6,6 +6,8 @@ import ( "runtime" "time" + "github.com/hashicorp/golang-lru/v2/expirable" + "github.com/matishsiao/goInfo" "github.com/shirou/gopsutil/cpu" "github.com/shirou/gopsutil/disk" "github.com/shirou/gopsutil/mem" @@ -52,11 +54,20 @@ type DiskStatus struct { // NewDiskStatus returns a new disk status of the server func NewDiskStatus() DiskStatus { - diskStat, _ := disk.Usage("/") - return DiskStatus{ - Total: diskStat.Total, - Used: diskStat.Used, - } + cache := expirable.NewLRU[string, DiskStatus](1, nil, 30*time.Minute) + return func() DiskStatus { + r, ok := cache.Get("disk_status") + if ok { + return r + } + diskStat, _ := disk.Usage("/") + status := DiskStatus{ + Total: diskStat.Total, + Used: diskStat.Used, + } + cache.Add("disk_status", status) + return status + }() } // MemoryStatus represents the memory status of the server @@ -76,31 +87,44 @@ func NewMemoryStatus() MemoryStatus { // GoStatus represents the Go status of the server type GoStatus struct { - NumGoroutines int `json:"num_goroutines"` - NumCgoCalls int64 `json:"num_cgo_calls"` + Version string `json:"version"` + NumGoroutines int `json:"num_goroutines"` + NumCgoCalls int64 `json:"num_cgo_calls"` + MemoryUsage int64 `json:"memory_usage"` } // NewGoStatus returns a new Go status of the server func NewGoStatus() GoStatus { + var m runtime.MemStats + runtime.ReadMemStats(&m) numGoroutines := runtime.NumGoroutine() numCgoCalls := runtime.NumCgoCall() return GoStatus{ NumGoroutines: numGoroutines, + Version: runtime.Version(), NumCgoCalls: numCgoCalls, + MemoryUsage: int64(m.Alloc), } } // OSStatus represents the OS status of the server type OSStatus struct { - OS string `json:"os"` - Arch string `json:"arch"` + Name string `json:"name"` + Version string `json:"version"` + Arch string `json:"arch"` } // NewOSStatus returns a new OS status of the server func NewOSStatus() OSStatus { + var version string = "unknown" + gi, err := goInfo.GetInfo() + if err == nil { + version = gi.Core + } return OSStatus{ - OS: runtime.GOOS, - Arch: runtime.GOARCH, + Name: runtime.GOOS, + Version: version, + Arch: runtime.GOARCH, } } @@ -133,6 +157,7 @@ type Status struct { DiskStatus `json:"disk,omitempty,omitzero"` MemoryStatus `json:"memory,omitempty,omitzero"` GoStatus `json:"go,omitempty,omitzero"` + UserStatus `json:"user,omitempty,omitzero"` Timestamp time.Time `json:"timestamp,omitzero"` } @@ -143,6 +168,7 @@ type StatusGrabber struct { withCPU bool withDisk bool withMemory bool + withUser bool withGo bool } @@ -192,6 +218,12 @@ func (s StatusGrabber) WithMemory() StatusGrabber { return s } +// WithUser adds user status to the status +func (s StatusGrabber) WithUser() StatusGrabber { + s.withUser = true + return s +} + // WithGo adds Go status to the status func (s StatusGrabber) WithGo() StatusGrabber { s.withGo = true @@ -219,5 +251,8 @@ func (s StatusGrabber) Grab() Status { if s.withGo { status.GoStatus = NewGoStatus() } + if s.withUser { + status.UserStatus = NewUserStatus() + } return status }