lixen/logwisp
SHA256
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

v0.3.3 pipeline rate limiter added

Browse Source
This commit is contained in:
Lixen Wraith
2025-07-13 03:20:47 -04:00
parent 0accb5f2d3
commit cc27f5cc1c
17 changed files with 742 additions and 588 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
src/internal/sink/http.go
View File

@ -12,7 +12,7 @@ import (
"time"
"logwisp/src/internal/config"
"logwisp/src/internal/ratelimit"
"logwisp/src/internal/netlimit"
"logwisp/src/internal/source"
"logwisp/src/internal/version"
@ -40,8 +40,8 @@ type HTTPSink struct {
// For router integration
standalone bool
// Rate limiting
rateLimiter *ratelimit.Limiter
// Net limiting
netLimiter *netlimit.Limiter
// Statistics
totalProcessed atomic.Uint64
@ -56,7 +56,7 @@ type HTTPConfig struct {
StatusPath string
Heartbeat config.HeartbeatConfig
SSL *config.SSLConfig
RateLimit *config.RateLimitConfig
NetLimit *config.NetLimitConfig
}
// NewHTTPSink creates a new HTTP streaming sink
@ -95,30 +95,30 @@ func NewHTTPSink(options map[string]any, logger *log.Logger) (*HTTPSink, error)
}
}
// Extract rate limit config
if rl, ok := options["rate_limit"].(map[string]any); ok {
cfg.RateLimit = &config.RateLimitConfig{}
cfg.RateLimit.Enabled, _ = rl["enabled"].(bool)
// Extract net limit config
if rl, ok := options["net_limit"].(map[string]any); ok {
cfg.NetLimit = &config.NetLimitConfig{}
cfg.NetLimit.Enabled, _ = rl["enabled"].(bool)
if rps, ok := toFloat(rl["requests_per_second"]); ok {
cfg.RateLimit.RequestsPerSecond = rps
cfg.NetLimit.RequestsPerSecond = rps
}
if burst, ok := toInt(rl["burst_size"]); ok {
cfg.RateLimit.BurstSize = burst
cfg.NetLimit.BurstSize = burst
}
if limitBy, ok := rl["limit_by"].(string); ok {
cfg.RateLimit.LimitBy = limitBy
cfg.NetLimit.LimitBy = limitBy
}
if respCode, ok := toInt(rl["response_code"]); ok {
cfg.RateLimit.ResponseCode = respCode
cfg.NetLimit.ResponseCode = respCode
}
if msg, ok := rl["response_message"].(string); ok {
cfg.RateLimit.ResponseMessage = msg
cfg.NetLimit.ResponseMessage = msg
}
if maxPerIP, ok := toInt(rl["max_connections_per_ip"]); ok {
cfg.RateLimit.MaxConnectionsPerIP = maxPerIP
cfg.NetLimit.MaxConnectionsPerIP = maxPerIP
}
if maxTotal, ok := toInt(rl["max_total_connections"]); ok {
cfg.RateLimit.MaxTotalConnections = maxTotal
cfg.NetLimit.MaxTotalConnections = maxTotal
}
}
@ -134,9 +134,9 @@ func NewHTTPSink(options map[string]any, logger *log.Logger) (*HTTPSink, error)
}
h.lastProcessed.Store(time.Time{})
// Initialize rate limiter if configured
if cfg.RateLimit != nil && cfg.RateLimit.Enabled {
h.rateLimiter = ratelimit.New(*cfg.RateLimit, logger)
// Initialize net limiter if configured
if cfg.NetLimit != nil && cfg.NetLimit.Enabled {
h.netLimiter = netlimit.New(*cfg.NetLimit, logger)
}
return h, nil
@ -212,9 +212,9 @@ func (h *HTTPSink) Stop() {
func (h *HTTPSink) GetStats() SinkStats {
lastProc, _ := h.lastProcessed.Load().(time.Time)
var rateLimitStats map[string]any
if h.rateLimiter != nil {
rateLimitStats = h.rateLimiter.GetStats()
var netLimitStats map[string]any
if h.netLimiter != nil {
netLimitStats = h.netLimiter.GetStats()
}
return SinkStats{
@ -230,7 +230,7 @@ func (h *HTTPSink) GetStats() SinkStats {
"stream": h.streamPath,
"status": h.statusPath,
},
"rate_limit": rateLimitStats,
"net_limit": netLimitStats,
},
}
}
@ -248,9 +248,9 @@ func (h *HTTPSink) RouteRequest(ctx *fasthttp.RequestCtx) {
}
func (h *HTTPSink) requestHandler(ctx *fasthttp.RequestCtx) {
// Check rate limit first
// Check net limit first
remoteAddr := ctx.RemoteAddr().String()
if allowed, statusCode, message := h.rateLimiter.CheckHTTP(remoteAddr); !allowed {
if allowed, statusCode, message := h.netLimiter.CheckHTTP(remoteAddr); !allowed {
ctx.SetStatusCode(statusCode)
ctx.SetContentType("application/json")
json.NewEncoder(ctx).Encode(map[string]any{
@ -279,11 +279,11 @@ func (h *HTTPSink) requestHandler(ctx *fasthttp.RequestCtx) {
}
func (h *HTTPSink) handleStream(ctx *fasthttp.RequestCtx) {
// Track connection for rate limiting
// Track connection for net limiting
remoteAddr := ctx.RemoteAddr().String()
if h.rateLimiter != nil {
h.rateLimiter.AddConnection(remoteAddr)
defer h.rateLimiter.RemoveConnection(remoteAddr)
if h.netLimiter != nil {
h.netLimiter.AddConnection(remoteAddr)
defer h.netLimiter.RemoveConnection(remoteAddr)
}
// Set SSE headers
@ -450,11 +450,11 @@ func (h *HTTPSink) formatHeartbeat() string {
func (h *HTTPSink) handleStatus(ctx *fasthttp.RequestCtx) {
ctx.SetContentType("application/json")
var rateLimitStats any
if h.rateLimiter != nil {
rateLimitStats = h.rateLimiter.GetStats()
var netLimitStats any
if h.netLimiter != nil {
netLimitStats = h.netLimiter.GetStats()
} else {
rateLimitStats = map[string]any{
netLimitStats = map[string]any{
"enabled": false,
}
}
@ -483,7 +483,7 @@ func (h *HTTPSink) handleStatus(ctx *fasthttp.RequestCtx) {
"ssl": map[string]bool{
"enabled": h.config.SSL != nil && h.config.SSL.Enabled,
},
"rate_limit": rateLimitStats,
"net_limit": netLimitStats,
},
}

56
src/internal/sink/tcp.go
View File

@ -11,7 +11,7 @@ import (
"time"
"logwisp/src/internal/config"
"logwisp/src/internal/ratelimit"
"logwisp/src/internal/netlimit"
"logwisp/src/internal/source"
"github.com/lixenwraith/log"
@ -29,7 +29,7 @@ type TCPSink struct {
engine *gnet.Engine
engineMu sync.Mutex
wg sync.WaitGroup
rateLimiter *ratelimit.Limiter
netLimiter *netlimit.Limiter
logger *log.Logger
// Statistics
@ -43,7 +43,7 @@ type TCPConfig struct {
BufferSize int
Heartbeat config.HeartbeatConfig
SSL *config.SSLConfig
RateLimit *config.RateLimitConfig
NetLimit *config.NetLimitConfig
}
// NewTCPSink creates a new TCP streaming sink
@ -74,30 +74,30 @@ func NewTCPSink(options map[string]any, logger *log.Logger) (*TCPSink, error) {
}
}
// Extract rate limit config
if rl, ok := options["rate_limit"].(map[string]any); ok {
cfg.RateLimit = &config.RateLimitConfig{}
cfg.RateLimit.Enabled, _ = rl["enabled"].(bool)
// Extract net limit config
if rl, ok := options["net_limit"].(map[string]any); ok {
cfg.NetLimit = &config.NetLimitConfig{}
cfg.NetLimit.Enabled, _ = rl["enabled"].(bool)
if rps, ok := toFloat(rl["requests_per_second"]); ok {
cfg.RateLimit.RequestsPerSecond = rps
cfg.NetLimit.RequestsPerSecond = rps
}
if burst, ok := toInt(rl["burst_size"]); ok {
cfg.RateLimit.BurstSize = burst
cfg.NetLimit.BurstSize = burst
}
if limitBy, ok := rl["limit_by"].(string); ok {
cfg.RateLimit.LimitBy = limitBy
cfg.NetLimit.LimitBy = limitBy
}
if respCode, ok := toInt(rl["response_code"]); ok {
cfg.RateLimit.ResponseCode = respCode
cfg.NetLimit.ResponseCode = respCode
}
if msg, ok := rl["response_message"].(string); ok {
cfg.RateLimit.ResponseMessage = msg
cfg.NetLimit.ResponseMessage = msg
}
if maxPerIP, ok := toInt(rl["max_connections_per_ip"]); ok {
cfg.RateLimit.MaxConnectionsPerIP = maxPerIP
cfg.NetLimit.MaxConnectionsPerIP = maxPerIP
}
if maxTotal, ok := toInt(rl["max_total_connections"]); ok {
cfg.RateLimit.MaxTotalConnections = maxTotal
cfg.NetLimit.MaxTotalConnections = maxTotal
}
}
@ -110,8 +110,8 @@ func NewTCPSink(options map[string]any, logger *log.Logger) (*TCPSink, error) {
}
t.lastProcessed.Store(time.Time{})
if cfg.RateLimit != nil && cfg.RateLimit.Enabled {
t.rateLimiter = ratelimit.New(*cfg.RateLimit, logger)
if cfg.NetLimit != nil && cfg.NetLimit.Enabled {
t.netLimiter = netlimit.New(*cfg.NetLimit, logger)
}
return t, nil
@ -194,9 +194,9 @@ func (t *TCPSink) Stop() {
func (t *TCPSink) GetStats() SinkStats {
lastProc, _ := t.lastProcessed.Load().(time.Time)
var rateLimitStats map[string]any
if t.rateLimiter != nil {
rateLimitStats = t.rateLimiter.GetStats()
var netLimitStats map[string]any
if t.netLimiter != nil {
netLimitStats = t.netLimiter.GetStats()
}
return SinkStats{
@ -208,7 +208,7 @@ func (t *TCPSink) GetStats() SinkStats {
Details: map[string]any{
"port": t.config.Port,
"buffer_size": t.config.BufferSize,
"rate_limit": rateLimitStats,
"net_limit": netLimitStats,
},
}
}
@ -313,8 +313,8 @@ func (s *tcpServer) OnOpen(c gnet.Conn) (out []byte, action gnet.Action) {
remoteAddr := c.RemoteAddr().String()
s.sink.logger.Debug("msg", "TCP connection attempt", "remote_addr", remoteAddr)
// Check rate limit
if s.sink.rateLimiter != nil {
// Check net limit
if s.sink.netLimiter != nil {
// Parse the remote address to get proper net.Addr
remoteStr := c.RemoteAddr().String()
tcpAddr, err := net.ResolveTCPAddr("tcp", remoteStr)
@ -325,15 +325,15 @@ func (s *tcpServer) OnOpen(c gnet.Conn) (out []byte, action gnet.Action) {
return nil, gnet.Close
}
if !s.sink.rateLimiter.CheckTCP(tcpAddr) {
s.sink.logger.Warn("msg", "TCP connection rate limited",
if !s.sink.netLimiter.CheckTCP(tcpAddr) {
s.sink.logger.Warn("msg", "TCP connection net limited",
"remote_addr", remoteAddr)
// Silently close connection when rate limited
// Silently close connection when net limited
return nil, gnet.Close
}
// Track connection
s.sink.rateLimiter.AddConnection(remoteStr)
s.sink.netLimiter.AddConnection(remoteStr)
}
s.connections.Store(c, struct{}{})
@ -352,8 +352,8 @@ func (s *tcpServer) OnClose(c gnet.Conn, err error) gnet.Action {
remoteAddr := c.RemoteAddr().String()
// Remove connection tracking
if s.sink.rateLimiter != nil {
s.sink.rateLimiter.RemoveConnection(c.RemoteAddr().String())
if s.sink.netLimiter != nil {
s.sink.netLimiter.RemoveConnection(c.RemoteAddr().String())
}
newCount := s.sink.activeConns.Add(-1)