Account for the root

[rbdr/mobius] / hotline / server.go

diff --git a/hotline/server.go b/hotline/server.go
index c4d8f7f429d22bf27bbd4cca397a56d9b712ce9c..a521a6bfecd22adc2a28c68ca1a4b6062ca3bbed 100644 (file)
--- a/hotline/server.go
+++ b/hotline/server.go
@@ -9,6 +9,7 @@ import (
        "errors"
        "fmt"
        "golang.org/x/text/encoding/charmap"
+       "golang.org/x/time/rate"
        "io"
        "log"
        "log/slog"
@@ -37,6 +38,8 @@ type Server struct {
        NetInterface string
        Port         int
 
+       rateLimiters map[string]*rate.Limiter
+
        handlers map[TranType]HandlerFunc
 
        Config Config
@@ -98,6 +101,7 @@ func NewServer(options ...Option) (*Server, error) {
        server := Server{
                handlers:        make(map[TranType]HandlerFunc),
                outbox:          make(chan Transaction),
+               rateLimiters:    make(map[string]*rate.Limiter),
                FS:              &OSFileStore{},
                ChatMgr:         NewMemChatManager(),
                ClientMgr:       NewMemClientMgr(),
@@ -170,7 +174,7 @@ func (s *Server) ServeFileTransfers(ctx context.Context, ln net.Listener) error
                        )
 
                        if err != nil {
-                               s.Logger.Error("file transfer error", "reason", err)
+                               s.Logger.Error("file transfer error", "err", err)
                        }
                }()
        }
@@ -202,6 +206,10 @@ func (s *Server) processOutbox() {
        }
 }
 
+// perIPRateLimit controls how frequently an IP address can connect before being throttled.
+// 0.5 = 1 connection every 2 seconds
+const perIPRateLimit = rate.Limit(0.5)
+
 func (s *Server) Serve(ctx context.Context, ln net.Listener) error {
        for {
                select {
@@ -216,13 +224,29 @@ func (s *Server) Serve(ctx context.Context, ln net.Listener) error {
                        }
 
                        go func() {
+                               ipAddr := strings.Split(conn.RemoteAddr().(*net.TCPAddr).String(), ":")[0]
+
                                connCtx := context.WithValue(ctx, contextKeyReq, requestCtx{
                                        remoteAddr: conn.RemoteAddr().String(),
                                })
 
-                               s.Logger.Info("Connection established", "addr", conn.RemoteAddr())
+                               s.Logger.Info("Connection established", "ip", ipAddr)
                                defer conn.Close()
 
+                               // Check if we have an existing rate limit for the IP and create one if we do not.
+                               rl, ok := s.rateLimiters[ipAddr]
+                               if !ok {
+                                       rl = rate.NewLimiter(perIPRateLimit, 1)
+                                       s.rateLimiters[ipAddr] = rl
+                               }
+
+                               // Check if the rate limit is exceeded and close the connection if so.
+                               if !rl.Allow() {
+                                       s.Logger.Info("Rate limit exceeded", "RemoteAddr", conn.RemoteAddr())
+                                       conn.Close()
+                                       return
+                               }
+
                                if err := s.handleNewConnection(connCtx, conn, conn.RemoteAddr().String()); err != nil {
                                        if err == io.EOF {
                                                s.Logger.Info("Client disconnected", "RemoteAddr", conn.RemoteAddr())
@@ -241,6 +265,10 @@ const trackerUpdateFrequency = 300
 // registerWithTrackers runs every trackerUpdateFrequency seconds to update the server's tracker entry on all configured
 // trackers.
 func (s *Server) registerWithTrackers(ctx context.Context) {
+       if s.Config.EnableTrackerRegistration {
+               s.Logger.Info("Tracker registration enabled", "trackers", s.Config.Trackers)
+       }
+
        for {
                if s.Config.EnableTrackerRegistration {
                        for _, t := range s.Config.Trackers {
@@ -400,7 +428,7 @@ func (s *Server) handleNewConnection(ctx context.Context, rwc io.ReadWriteCloser
                        return err
                }
 
-               c.Logger.Info("Login failed", "clientVersion", fmt.Sprintf("%x", c.Version))
+               c.Logger.Info("Incorrect login")
 
                return nil
        }
@@ -575,10 +603,20 @@ func (s *Server) handleFileTransfer(ctx context.Context, rwc io.ReadWriter) erro
                        s.Stats.Decrement(StatUploadsInProgress)
                }()
 
+               var transferSizeValue uint32
+               switch len(fileTransfer.TransferSize) {
+                       case 2: // 16-bit
+                               transferSizeValue = uint32(binary.BigEndian.Uint16(fileTransfer.TransferSize))
+                       case 4: // 32-bit
+                               transferSizeValue = binary.BigEndian.Uint32(fileTransfer.TransferSize)
+                       default:
+                               rLogger.Warn("Unexpected TransferSize length: %d bytes", len(fileTransfer.TransferSize))
+               }
+
                rLogger.Info(
                        "Folder upload started",
                        "dstPath", fullPath,
-                       "TransferSize", binary.BigEndian.Uint32(fileTransfer.TransferSize),
+                       "TransferSize", transferSizeValue,
                        "FolderItemCount", fileTransfer.FolderItemCount,
                )