Skip to content

Commit a97ecb4

Browse files
0x46616c6bOpenCode
0x46616c6b
and
OpenCode
committed
Use named loggers to distinguish lookup and policy server in logs
Add a *zap.Logger field to LookupServer, PolicyServer, and TCPServerConfig using zap's Named() method. This adds a 'logger' field to all JSON log output, making it easy to filter and distinguish between the two servers. Co-Authored-By: OpenCode <noreply@opencode.ai>
1 parent 05f82e7 commit a97ecb4

6 files changed

Lines changed: 66 additions & 59 deletions

File tree

‎lookup.go‎

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -32,18 +32,20 @@ func (r *SocketmapResponse) String() string {
3232
// It implements the ConnectionHandler interface.
3333
type LookupServer struct {
3434
client UserliService
35+
logger *zap.Logger
3536
}
3637

3738
// NewLookupServer creates a new LookupServer with the given UserliService
38-
func NewLookupServer(client UserliService) *LookupServer {
39-
return &LookupServer{client: client}
39+
func NewLookupServer(client UserliService, logger *zap.Logger) *LookupServer {
40+
return &LookupServer{client: client, logger: logger}
4041
}
4142

4243
// StartLookupServer starts the lookup server on the given address
4344
func StartLookupServer(ctx context.Context, wg *sync.WaitGroup, addr string, server *LookupServer) {
4445
config := TCPServerConfig{
45-
Name: "socketmap",
46-
Addr: addr,
46+
Name: "lookup",
47+
Addr: addr,
48+
Logger: server.logger,
4749
OnConnectionAcquired: func() {
4850
activeConnections.Inc()
4951
},
@@ -81,7 +83,7 @@ func (s *LookupServer) HandleConnection(ctx context.Context, conn net.Conn) {
8183
if errors.As(err, &netErr) && netErr.Timeout() {
8284
return
8385
}
84-
logger.Debug("Failed to decode request", zap.Error(err))
86+
s.logger.Debug("Failed to decode request", zap.Error(err))
8587
return
8688
}
8789
request := string(requestBytes)
@@ -91,7 +93,7 @@ func (s *LookupServer) HandleConnection(ctx context.Context, conn net.Conn) {
9193
// Parse the request: "name key"
9294
parts := strings.SplitN(strings.TrimSpace(request), " ", 2)
9395
if len(parts) != 2 {
94-
logger.Error("Invalid request format", zap.String("request", request))
96+
s.logger.Error("Invalid request format", zap.String("request", request))
9597
response := &SocketmapResponse{Status: "PERM", Data: "Invalid request format"}
9698
s.writeResponse(encoder, conn, response, now, "invalid")
9799
continue
@@ -100,7 +102,7 @@ func (s *LookupServer) HandleConnection(ctx context.Context, conn net.Conn) {
100102
mapName := parts[0]
101103
key := parts[1]
102104

103-
logger.Debug("Processing socketmap request",
105+
s.logger.Debug("Processing socketmap request",
104106
zap.String("map", mapName),
105107
zap.String("key", key))
106108

@@ -125,7 +127,7 @@ func (s *LookupServer) processRequest(ctx context.Context, mapName, key string)
125127
case "senders":
126128
return s.handleSenders(reqCtx, key)
127129
default:
128-
logger.Error("Unknown map name", zap.String("map", mapName))
130+
s.logger.Error("Unknown map name", zap.String("map", mapName))
129131
return &SocketmapResponse{Status: "PERM", Data: "Unknown map name"}
130132
}
131133
}
@@ -134,7 +136,7 @@ func (s *LookupServer) processRequest(ctx context.Context, mapName, key string)
134136
func (s *LookupServer) handleAlias(ctx context.Context, key string) *SocketmapResponse {
135137
aliases, err := s.client.GetAliases(ctx, key)
136138
if err != nil {
137-
logger.Error("Error fetching aliases", zap.String("key", key), zap.Error(err))
139+
s.logger.Error("Error fetching aliases", zap.String("key", key), zap.Error(err))
138140
return &SocketmapResponse{Status: "TEMP", Data: "Error fetching aliases"}
139141
}
140142

@@ -149,7 +151,7 @@ func (s *LookupServer) handleAlias(ctx context.Context, key string) *SocketmapRe
149151
func (s *LookupServer) handleDomain(ctx context.Context, key string) *SocketmapResponse {
150152
exists, err := s.client.GetDomain(ctx, key)
151153
if err != nil {
152-
logger.Error("Error fetching domain", zap.String("key", key), zap.Error(err))
154+
s.logger.Error("Error fetching domain", zap.String("key", key), zap.Error(err))
153155
return &SocketmapResponse{Status: "TEMP", Data: "Error fetching domain"}
154156
}
155157

@@ -164,7 +166,7 @@ func (s *LookupServer) handleDomain(ctx context.Context, key string) *SocketmapR
164166
func (s *LookupServer) handleMailbox(ctx context.Context, key string) *SocketmapResponse {
165167
exists, err := s.client.GetMailbox(ctx, key)
166168
if err != nil {
167-
logger.Error("Error fetching mailbox", zap.String("key", key), zap.Error(err))
169+
s.logger.Error("Error fetching mailbox", zap.String("key", key), zap.Error(err))
168170
return &SocketmapResponse{Status: "TEMP", Data: "Error fetching mailbox"}
169171
}
170172

@@ -179,7 +181,7 @@ func (s *LookupServer) handleMailbox(ctx context.Context, key string) *Socketmap
179181
func (s *LookupServer) handleSenders(ctx context.Context, key string) *SocketmapResponse {
180182
senders, err := s.client.GetSenders(ctx, key)
181183
if err != nil {
182-
logger.Error("Error fetching senders", zap.String("key", key), zap.Error(err))
184+
s.logger.Error("Error fetching senders", zap.String("key", key), zap.Error(err))
183185
return &SocketmapResponse{Status: "TEMP", Data: "Error fetching senders"}
184186
}
185187

@@ -202,7 +204,7 @@ func (s *LookupServer) writeResponse(encoder *netstring.Encoder, conn net.Conn,
202204
status = "error"
203205
}
204206

205-
logger.Debug("Writing socketmap response",
207+
s.logger.Debug("Writing socketmap response",
206208
zap.String("response", response.String()),
207209
zap.String("map", mapName),
208210
zap.String("status", status))
@@ -213,7 +215,7 @@ func (s *LookupServer) writeResponse(encoder *netstring.Encoder, conn net.Conn,
213215
// Encode and send the response
214216
err := encoder.EncodeString(netstring.NoKey, response.String())
215217
if err != nil {
216-
logger.Error("Error writing response",
218+
s.logger.Error("Error writing response",
217219
zap.String("response", response.String()),
218220
zap.String("map", mapName),
219221
zap.String("status", status),

‎lookup_test.go‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ func (s *ServerTestSuite) TearDownTest() {
4949
// TestStartLookupServer_BasicFunctionality tests basic server startup and shutdown
5050
func (s *ServerTestSuite) TestStartLookupServer_BasicFunctionality() {
5151
mock := &MockUserliService{}
52-
server := NewLookupServer(mock)
52+
server := NewLookupServer(mock, zap.NewNop())
5353

5454
ctx, cancel := context.WithCancel(context.Background())
5555
var wg sync.WaitGroup
@@ -84,7 +84,7 @@ func (s *ServerTestSuite) TestStartLookupServer_BasicFunctionality() {
8484
// TestStartLookupServer_InvalidAddress tests server behavior with invalid address
8585
func (s *ServerTestSuite) TestStartLookupServer_InvalidAddress() {
8686
mock := &MockUserliService{}
87-
server := NewLookupServer(mock)
87+
server := NewLookupServer(mock, zap.NewNop())
8888

8989
ctx, cancel := context.WithCancel(context.Background())
9090
defer cancel()
@@ -116,7 +116,7 @@ func (s *ServerTestSuite) TestStartLookupServer_ConnectionHandling() {
116116
mockService := &MockUserliService{}
117117
// Mock a successful domain lookup
118118
mockService.On("GetDomain", mock.Anything, "example.com").Return(true, nil)
119-
server := NewLookupServer(mockService)
119+
server := NewLookupServer(mockService, zap.NewNop())
120120

121121
ctx, cancel := context.WithCancel(context.Background())
122122
defer cancel()
@@ -155,7 +155,7 @@ func (s *ServerTestSuite) TestStartLookupServer_ConnectionHandling() {
155155
// TestStartLookupServer_GracefulShutdown tests graceful shutdown with active connections
156156
func (s *ServerTestSuite) TestStartLookupServer_GracefulShutdown() {
157157
mockService := &MockUserliService{}
158-
server := NewLookupServer(mockService)
158+
server := NewLookupServer(mockService, zap.NewNop())
159159

160160
ctx, cancel := context.WithCancel(context.Background())
161161
var wg sync.WaitGroup
@@ -207,7 +207,7 @@ func (s *ServerTestSuite) TestStartLookupServer_GracefulShutdown() {
207207
func (s *ServerTestSuite) TestHandleLookupConnection() {
208208
mockService := &MockUserliService{}
209209
mockService.On("GetDomain", mock.Anything, "example.com").Return(true, nil)
210-
server := NewLookupServer(mockService)
210+
server := NewLookupServer(mockService, zap.NewNop())
211211

212212
// Create a pipe to simulate a connection
213213
serverConn, client := net.Pipe()
@@ -233,7 +233,7 @@ func (s *ServerTestSuite) TestHandleLookupConnection() {
233233
// TestStartLookupServer_ConnectionPoolLimit tests connection pool limits
234234
func (s *ServerTestSuite) TestStartLookupServer_ConnectionPoolLimit() {
235235
mockService := &MockUserliService{}
236-
server := NewLookupServer(mockService)
236+
server := NewLookupServer(mockService, zap.NewNop())
237237

238238
ctx, cancel := context.WithCancel(context.Background())
239239
defer cancel()
@@ -280,7 +280,7 @@ func (s *ServerTestSuite) TestHandleLookupConnection_MultipleRequests() {
280280
mockService := &MockUserliService{}
281281
mockService.On("GetDomain", mock.Anything, "example.com").Return(true, nil)
282282
mockService.On("GetDomain", mock.Anything, "example.org").Return(false, nil)
283-
server := NewLookupServer(mockService)
283+
server := NewLookupServer(mockService, zap.NewNop())
284284

285285
serverConn, client := net.Pipe()
286286
defer serverConn.Close()
@@ -318,7 +318,7 @@ func (s *ServerTestSuite) TestHandleLookupConnection_AliasLookup() {
318318
mockService.On("GetAliases", mock.Anything, "alias@example.com").Return([]string{"user1@example.com", "user2@example.com"}, nil)
319319
mockService.On("GetAliases", mock.Anything, "unknown@example.com").Return([]string{}, nil)
320320
mockService.On("GetAliases", mock.Anything, "error@example.com").Return([]string(nil), io.ErrUnexpectedEOF)
321-
server := NewLookupServer(mockService)
321+
server := NewLookupServer(mockService, zap.NewNop())
322322

323323
serverConn, client := net.Pipe()
324324
defer serverConn.Close()
@@ -356,7 +356,7 @@ func (s *ServerTestSuite) TestHandleLookupConnection_MailboxLookup() {
356356
mockService.On("GetMailbox", mock.Anything, "user@example.com").Return(true, nil)
357357
mockService.On("GetMailbox", mock.Anything, "unknown@example.com").Return(false, nil)
358358
mockService.On("GetMailbox", mock.Anything, "error@example.com").Return(false, io.ErrUnexpectedEOF)
359-
server := NewLookupServer(mockService)
359+
server := NewLookupServer(mockService, zap.NewNop())
360360

361361
serverConn, client := net.Pipe()
362362
defer serverConn.Close()
@@ -394,7 +394,7 @@ func (s *ServerTestSuite) TestHandleLookupConnection_SendersLookup() {
394394
mockService.On("GetSenders", mock.Anything, "user@example.com").Return([]string{"alias1@example.com", "alias2@example.com"}, nil)
395395
mockService.On("GetSenders", mock.Anything, "unknown@example.com").Return([]string{}, nil)
396396
mockService.On("GetSenders", mock.Anything, "error@example.com").Return([]string(nil), io.ErrUnexpectedEOF)
397-
server := NewLookupServer(mockService)
397+
server := NewLookupServer(mockService, zap.NewNop())
398398

399399
serverConn, client := net.Pipe()
400400
defer serverConn.Close()
@@ -430,7 +430,7 @@ func (s *ServerTestSuite) TestHandleLookupConnection_SendersLookup() {
430430
func (s *ServerTestSuite) TestHandleLookupConnection_DomainError() {
431431
mockService := &MockUserliService{}
432432
mockService.On("GetDomain", mock.Anything, "error.com").Return(false, io.ErrUnexpectedEOF)
433-
server := NewLookupServer(mockService)
433+
server := NewLookupServer(mockService, zap.NewNop())
434434

435435
serverConn, client := net.Pipe()
436436
defer serverConn.Close()
@@ -450,7 +450,7 @@ func (s *ServerTestSuite) TestHandleLookupConnection_DomainError() {
450450
// TestHandleLookupConnection_UnknownMap tests unknown map name handling
451451
func (s *ServerTestSuite) TestHandleLookupConnection_UnknownMap() {
452452
mockService := &MockUserliService{}
453-
server := NewLookupServer(mockService)
453+
server := NewLookupServer(mockService, zap.NewNop())
454454

455455
serverConn, client := net.Pipe()
456456
defer serverConn.Close()
@@ -469,7 +469,7 @@ func (s *ServerTestSuite) TestHandleLookupConnection_UnknownMap() {
469469
// TestHandleLookupConnection_InvalidFormat tests invalid request format
470470
func (s *ServerTestSuite) TestHandleLookupConnection_InvalidFormat() {
471471
mockService := &MockUserliService{}
472-
server := NewLookupServer(mockService)
472+
server := NewLookupServer(mockService, zap.NewNop())
473473

474474
serverConn, client := net.Pipe()
475475
defer serverConn.Close()
@@ -489,7 +489,7 @@ func (s *ServerTestSuite) TestHandleLookupConnection_InvalidFormat() {
489489
// TestHandleLookupConnection_ContextCancelled tests context cancellation
490490
func (s *ServerTestSuite) TestHandleLookupConnection_ContextCancelled() {
491491
mockService := &MockUserliService{}
492-
server := NewLookupServer(mockService)
492+
server := NewLookupServer(mockService, zap.NewNop())
493493

494494
serverConn, client := net.Pipe()
495495

‎main.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,13 +42,13 @@ func main() {
4242
}
4343

4444
userli := NewUserli(config.UserliToken, config.UserliBaseURL, WithDelimiter(config.PostfixRecipientDelimiter))
45-
lookupServer := NewLookupServer(userli)
45+
lookupServer := NewLookupServer(userli, logger.Named("lookup"))
4646

4747
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
4848
defer stop()
4949

5050
rateLimiter := NewRateLimiter(ctx)
51-
policyServer := NewPolicyServer(userli, rateLimiter)
51+
policyServer := NewPolicyServer(userli, rateLimiter, logger.Named("policy"))
5252

5353
var wg sync.WaitGroup
5454

‎policy.go‎

Lines changed: 14 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -18,21 +18,24 @@ import (
1818
type PolicyServer struct {
1919
client UserliService
2020
rateLimiter *RateLimiter
21+
logger *zap.Logger
2122
}
2223

2324
// NewPolicyServer creates a new PolicyServer with the given UserliService
24-
func NewPolicyServer(client UserliService, rateLimiter *RateLimiter) *PolicyServer {
25+
func NewPolicyServer(client UserliService, rateLimiter *RateLimiter, logger *zap.Logger) *PolicyServer {
2526
return &PolicyServer{
2627
client: client,
2728
rateLimiter: rateLimiter,
29+
logger: logger,
2830
}
2931
}
3032

3133
// StartPolicyServer starts the policy server on the given address
3234
func StartPolicyServer(ctx context.Context, wg *sync.WaitGroup, addr string, server *PolicyServer) {
3335
config := TCPServerConfig{
34-
Name: "policy",
35-
Addr: addr,
36+
Name: "policy",
37+
Addr: addr,
38+
Logger: server.logger,
3639
OnConnectionAcquired: func() {
3740
policyActiveConnections.Inc()
3841
},
@@ -65,15 +68,15 @@ func (p *PolicyServer) HandleConnection(ctx context.Context, conn net.Conn) {
6568
if errors.As(err, &netErr) && netErr.Timeout() {
6669
return
6770
}
68-
logger.Debug("Failed to read policy request", zap.Error(err))
71+
p.logger.Debug("Failed to read policy request", zap.Error(err))
6972
return
7073
}
7174

7275
response := p.handleRequest(ctx, request)
7376

7477
_ = conn.SetWriteDeadline(time.Now().Add(WriteTimeout))
7578
if err := p.writeResponse(conn, response); err != nil {
76-
logger.Error("Failed to write policy response", zap.Error(err))
79+
p.logger.Error("Failed to write policy response", zap.Error(err))
7780
return
7881
}
7982
}
@@ -162,7 +165,7 @@ func (p *PolicyServer) readRequest(reader *bufio.Reader) (*PolicyRequest, error)
162165
func (p *PolicyServer) handleRequest(ctx context.Context, req *PolicyRequest) string {
163166
startTime := time.Now()
164167

165-
logger.Debug("Processing policy request",
168+
p.logger.Debug("Processing policy request",
166169
zap.String("queue_id", req.QueueID),
167170
zap.String("sender", req.Sender),
168171
zap.String("sasl_username", req.SaslUsername),
@@ -184,7 +187,7 @@ func (p *PolicyServer) handleRequest(ctx context.Context, req *PolicyRequest) st
184187
}
185188

186189
if sender == "" {
187-
logger.Debug("No sender identity found, allowing message")
190+
p.logger.Debug("No sender identity found, allowing message")
188191
policyRequestsTotal.WithLabelValues("check", "dunno").Inc()
189192
policyRequestDuration.WithLabelValues("check", "dunno").Observe(time.Since(startTime).Seconds())
190193
return "DUNNO"
@@ -197,7 +200,7 @@ func (p *PolicyServer) handleRequest(ctx context.Context, req *PolicyRequest) st
197200
quota, err := p.client.GetQuota(quotaCtx, sender)
198201
if err != nil {
199202
// API error - fail open (allow the message)
200-
logger.Warn("Failed to fetch quota, allowing message",
203+
p.logger.Warn("Failed to fetch quota, allowing message",
201204
zap.String("queue_id", req.QueueID),
202205
zap.String("sender", sender), zap.Error(err))
203206
policyRequestsTotal.WithLabelValues("check", "error").Inc()
@@ -207,7 +210,7 @@ func (p *PolicyServer) handleRequest(ctx context.Context, req *PolicyRequest) st
207210

208211
// No limits configured (both 0 means unlimited)
209212
if quota.PerHour == 0 && quota.PerDay == 0 {
210-
logger.Debug("No quota limits configured", zap.String("sender", sender))
213+
p.logger.Debug("No quota limits configured", zap.String("sender", sender))
211214
policyRequestsTotal.WithLabelValues("check", "dunno").Inc()
212215
policyRequestDuration.WithLabelValues("check", "dunno").Observe(time.Since(startTime).Seconds())
213216
return "DUNNO"
@@ -220,7 +223,7 @@ func (p *PolicyServer) handleRequest(ctx context.Context, req *PolicyRequest) st
220223
quotaChecksTotal.WithLabelValues("checked").Inc()
221224

222225
if !allowed {
223-
logger.Info("Rate limit exceeded",
226+
p.logger.Info("Rate limit exceeded",
224227
zap.String("queue_id", req.QueueID),
225228
zap.String("sender", sender),
226229
zap.Int("hour_count", hourCount),
@@ -235,7 +238,7 @@ func (p *PolicyServer) handleRequest(ctx context.Context, req *PolicyRequest) st
235238
return "REJECT Rate limit exceeded, please try again later"
236239
}
237240

238-
logger.Debug("Message allowed",
241+
p.logger.Debug("Message allowed",
239242
zap.String("queue_id", req.QueueID),
240243
zap.String("sender", sender),
241244
zap.Int("hour_count", hourCount),

0 commit comments

Comments
 (0)