Skip to content

Commit 11be04c

Browse files
tooryxcopybara-github
tooryx
authored and
copybara-github
committed
Extending the context with service information now takes such service rather than a port.
PiperOrigin-RevId: 871239254
1 parent a133b03 commit 11be04c

File tree

4 files changed

+41
-16
lines changed

4 files changed

+41
-16
lines changed

‎common/clients/llm/tools/httpclient/httpclient.go‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -138,9 +138,8 @@ func (h *Tool) numberOfRequests() int {
138138

139139
// Do performs an HTTP request against the service.
140140
func (h *Tool) Do(toolctx tool.Context, toolreq *Request) (*Response, error) {
141-
port := int(h.service.GetNetworkEndpoint().GetPort().GetPortNumber())
142141
uri := toolreq.URI
143-
ctx := log.ContextForModuleAndService(context.Background(), "clients/llm/httpclient", port)
142+
ctx := log.ContextForModuleAndService(context.Background(), "clients/llm/httpclient", h.service)
144143
ctx, cancel := context.WithTimeout(ctx, h.coreConfig.TimeoutPerRequest())
145144
defer cancel()
146145

@@ -168,26 +167,25 @@ func (h *Tool) Do(toolctx tool.Context, toolreq *Request) (*Response, error) {
168167
}
169168

170169
func (h *Tool) prepareRequest(ctx context.Context, toolreq *Request, path string) (*http.Request, error) {
171-
port := h.service.GetNetworkEndpoint().GetPort().GetPortNumber()
172170
uri := toolreq.URI
173171

174172
if !strings.HasPrefix(uri, "/") || strings.Contains(uri, "://") || strings.Contains(uri, ":") {
175173
return nil, ErrInvalidURI
176174
}
177175

178176
if !slices.Contains(h.config.GetAllowedMethods(), toolreq.Method) {
179-
log.DebugContextf(ctx, log.DebugLevelRequest, "%s invalid method for %q", port, toolreq.Method, uri)
177+
log.DebugContextf(ctx, log.DebugLevelRequest, "invalid method for %q: %q", toolreq.Method, uri)
180178
return nil, ErrInvalidMethod
181179
}
182180

183181
if h.numberOfRequests() >= int(h.config.GetMaxRequestsPerService()) {
184-
log.DebugContextf(ctx, log.DebugLevelRequest, "%s too many requests for %q", port, toolreq.Method, uri)
182+
log.DebugContextf(ctx, log.DebugLevelRequest, "too many requests for %q: %q", toolreq.Method, uri)
185183
return nil, ErrTooManyRequests
186184
}
187185

188186
for _, path := range h.badPaths {
189187
if path.MatchString(uri) {
190-
log.DebugContextf(ctx, log.DebugLevelRequest, "%s %q denied by regexp", port, toolreq.Method, uri)
188+
log.DebugContextf(ctx, log.DebugLevelRequest, "%q denied by regexp: %q", toolreq.Method, uri)
191189
return nil, ErrContentDenied
192190
}
193191
}

‎core/log/log.go‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@ import (
2222
"context"
2323
"fmt"
2424
"log"
25+
26+
nspb "github.com/google/tsunami-security-scanner/proto/go/network_service_go_proto"
2527
)
2628

2729
// Logger is Goonami's logging interface.
@@ -68,14 +70,15 @@ func ContextForModule(ctx context.Context, name string) context.Context {
6870
}
6971

7072
// ContextForService returns a new context with the service information (port) attached.
71-
func ContextForService(ctx context.Context, port int) context.Context {
73+
func ContextForService(ctx context.Context, service *nspb.NetworkService) context.Context {
74+
port := int(service.GetNetworkEndpoint().GetPort().GetPortNumber())
7275
return context.WithValue(ctx, serviceKey, port)
7376
}
7477

7578
// ContextForModuleAndService returns a new context with both the module name and the service
7679
// information (port) attached.
77-
func ContextForModuleAndService(ctx context.Context, name string, port int) context.Context {
78-
return ContextForModule(ContextForService(ctx, port), name)
80+
func ContextForModuleAndService(ctx context.Context, name string, service *nspb.NetworkService) context.Context {
81+
return ContextForModule(ContextForService(ctx, service), name)
7982
}
8083

8184
var logger Logger = &DefaultLogger{}

‎core/log/log_test.go‎

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,9 @@ import (
2222
"log"
2323
"os"
2424
"testing"
25+
26+
npb "github.com/google/tsunami-security-scanner/proto/go/network_go_proto"
27+
nspb "github.com/google/tsunami-security-scanner/proto/go/network_service_go_proto"
2528
)
2629

2730
var buf bytes.Buffer
@@ -133,6 +136,22 @@ func TestDefaultLoggerDebug(t *testing.T) {
133136
}
134137

135138
func TestDefaultLoggerPrefix(t *testing.T) {
139+
svc80 := nspb.NetworkService_builder{
140+
NetworkEndpoint: npb.NetworkEndpoint_builder{
141+
Port: npb.Port_builder{
142+
PortNumber: 80,
143+
}.Build(),
144+
}.Build(),
145+
}.Build()
146+
147+
svc443 := nspb.NetworkService_builder{
148+
NetworkEndpoint: npb.NetworkEndpoint_builder{
149+
Port: npb.Port_builder{
150+
PortNumber: 443,
151+
}.Build(),
152+
}.Build(),
153+
}.Build()
154+
136155
testCases := []struct {
137156
name string
138157
ctx context.Context
@@ -154,13 +173,13 @@ func TestDefaultLoggerPrefix(t *testing.T) {
154173
},
155174
{
156175
name: "when_service_metadata_service_prefix",
157-
ctx: ContextForService(context.Background(), 80),
176+
ctx: ContextForService(context.Background(), svc80),
158177
msg: "test",
159178
wantPrefix: "INFO [ 80 ] test",
160179
},
161180
{
162181
name: "when_both_metadata_both_prefix",
163-
ctx: ContextForModuleAndService(context.Background(), "my-module", 443),
182+
ctx: ContextForModuleAndService(context.Background(), "my-module", svc443),
164183
msg: "test",
165184
wantPrefix: "INFO [ 443 ] [ my-module ] test",
166185
},
@@ -192,7 +211,14 @@ func TestDefaultLoggerPrefix(t *testing.T) {
192211
}
193212

194213
func TestDefaultLoggerColors(t *testing.T) {
195-
ctx := ContextForModuleAndService(context.Background(), "my-module", 443)
214+
svc443 := nspb.NetworkService_builder{
215+
NetworkEndpoint: npb.NetworkEndpoint_builder{
216+
Port: npb.Port_builder{
217+
PortNumber: 443,
218+
}.Build(),
219+
}.Build(),
220+
}.Build()
221+
ctx := ContextForModuleAndService(context.Background(), "my-module", svc443)
196222
buf.Reset()
197223
l := &DefaultLogger{UseColors: true}
198224
l.InfoContext(ctx, "test message")

‎core/runner/simplerunner/simplerunner.go‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -236,7 +236,6 @@ func runConcurrent[E any](ctx context.Context, concurrency int, services []*nspb
236236
// coming directly from the port scan. But each fingerprinter can technically "split" that network
237237
// service into several ones.
238238
func (r *SimpleRunner) fingerprintService(ctx context.Context, svc *nspb.NetworkService) ([]*nspb.NetworkService, error) {
239-
port := int(svc.GetNetworkEndpoint().GetPort().GetPortNumber())
240239
var services []*nspb.NetworkService = []*nspb.NetworkService{svc}
241240
for _, fp := range r.fingerprinters {
242241
var accumulator []*nspb.NetworkService
@@ -246,7 +245,7 @@ func (r *SimpleRunner) fingerprintService(ctx context.Context, svc *nspb.Network
246245
return nil, ctx.Err()
247246
}
248247

249-
ctx = log.ContextForModuleAndService(ctx, fp.Name(), port)
248+
ctx = log.ContextForModuleAndService(ctx, fp.Name(), sv)
250249
res, err := fp.Fingerprint(ctx, sv)
251250
if err != nil {
252251
log.ErrorContextf(ctx, "fatal fingerprinting error: %s", err)
@@ -263,15 +262,14 @@ func (r *SimpleRunner) fingerprintService(ctx context.Context, svc *nspb.Network
263262
}
264263

265264
func (r *SimpleRunner) detectService(ctx context.Context, svc *nspb.NetworkService) ([]*dpb.DetectionReport, error) {
266-
port := int(svc.GetNetworkEndpoint().GetPort().GetPortNumber())
267265
var reports []*dpb.DetectionReport
268266

269267
for _, dt := range r.detectors {
270268
if ctx.Err() != nil {
271269
return nil, ctx.Err()
272270
}
273271

274-
ctx = log.ContextForModuleAndService(ctx, dt.Name(), port)
272+
ctx = log.ContextForModuleAndService(ctx, dt.Name(), svc)
275273
res, err := dt.Detect(ctx, svc)
276274
if err != nil {
277275
log.ErrorContextf(ctx, "fatal detection error: %s", err)

0 commit comments

Comments
 (0)