Adds MaxDiffEntries, MaxDatabaseEntries to config.

PiperOrigin-RevId: 639042423

api.go

@@ -41,8 +41,7 @@ const (
 
 // The api interface specifies wrappers around the Web Risk API.
 type api interface {
-	ListUpdate(ctx context.Context, threatType pb.ThreatType, versionToken []byte,
-		compressionTypes []pb.CompressionType) (*pb.ComputeThreatListDiffResponse, error)
+	ListUpdate(ctx context.Context, req *pb.ComputeThreatListDiffRequest) (*pb.ComputeThreatListDiffResponse, error)
 	HashLookup(ctx context.Context, hashPrefix []byte,
 		threatTypes []pb.ThreatType) (*pb.SearchHashesResponse, error)
 }
@@ -123,17 +122,16 @@ func (a *netAPI) parseError(httpResp *http.Response) error {
 }
 
 // ListUpdate issues a ComputeThreatListDiff API call and returns the response.
-func (a *netAPI) ListUpdate(ctx context.Context, threatType pb.ThreatType, versionToken []byte,
-	compressionTypes []pb.CompressionType) (*pb.ComputeThreatListDiffResponse, error) {
+func (a *netAPI) ListUpdate(ctx context.Context, req *pb.ComputeThreatListDiffRequest) (*pb.ComputeThreatListDiffResponse, error) {
 	resp := new(pb.ComputeThreatListDiffResponse)
 	u := *a.url // Make a copy of URL
 	// Add fields from ComputeThreatListDiffRequest to URL request
 	q := u.Query()
-	q.Set(threatTypeString, threatType.String())
-	if len(versionToken) != 0 {
-		q.Set(versionTokenString, base64.StdEncoding.EncodeToString(versionToken))
+	q.Set(threatTypeString, req.GetThreatType().String())
+	if len(req.GetVersionToken()) != 0 {
+		q.Set(versionTokenString, base64.StdEncoding.EncodeToString(req.GetVersionToken()))
 	}
-	for _, compressionType := range compressionTypes {
+	for _, compressionType := range req.GetConstraints().GetSupportedCompressions() {
 		q.Add(supportedCompressionsString, compressionType.String())
 	}
 	u.RawQuery = q.Encode()

api_test.go

@@ -39,9 +39,8 @@ type mockAPI struct {
 		threatTypes []pb.ThreatType) (*pb.SearchHashesResponse, error)
 }
 
-func (m *mockAPI) ListUpdate(ctx context.Context, threatType pb.ThreatType, versionToken []byte,
-	compressionTypes []pb.CompressionType) (*pb.ComputeThreatListDiffResponse, error) {
-	return m.listUpdate(ctx, threatType, versionToken, compressionTypes)
+func (m *mockAPI) ListUpdate(ctx context.Context, req *pb.ComputeThreatListDiffRequest) (*pb.ComputeThreatListDiffResponse, error) {
+	return m.listUpdate(ctx, req.GetThreatType(), req.GetVersionToken(), req.GetConstraints().GetSupportedCompressions())
 }
 
 func (m *mockAPI) HashLookup(ctx context.Context, hashPrefix []byte,
@@ -118,8 +117,14 @@ func TestNetAPI(t *testing.T) {
 			RawIndices: &pb.RawIndices{Indices: []int32{1, 2, 3}},
 		},
 	}
-	resp1, err := api.ListUpdate(context.Background(), wantReqThreatType, []byte{},
-		wantReqCompressionTypes)
+	req := &pb.ComputeThreatListDiffRequest{
+		ThreatType: wantReqThreatType,
+		Constraints: &pb.ComputeThreatListDiffRequest_Constraints{
+			SupportedCompressions: wantReqCompressionTypes,
+		},
+		VersionToken: []byte{},
+	}
+	resp1, err := api.ListUpdate(context.Background(), req)
 	gotResp = resp1
 	if err != nil {
 		t.Errorf("unexpected ListUpdate error: %v", err)

database.go

@@ -200,7 +200,10 @@ func (db *database) Update(ctx context.Context, api api) (time.Duration, bool) {
 		s = append(s, &pb.ComputeThreatListDiffRequest{
 			ThreatType: pb.ThreatType(td),
 			Constraints: &pb.ComputeThreatListDiffRequest_Constraints{
-				SupportedCompressions: db.config.compressionTypes},
+				SupportedCompressions: db.config.compressionTypes,
+				MaxDiffEntries:        db.config.MaxDiffEntries,
+				MaxDatabaseEntries:    db.config.MaxDatabaseEntries,
+			},
 			VersionToken: state,
 		})
 	}
@@ -212,7 +215,7 @@ func (db *database) Update(ctx context.Context, api api) (time.Duration, bool) {
 	last := db.config.now()
 	for _, req := range s {
 		// Query the API for the threat list and update the database.
-		resp, err := api.ListUpdate(ctx, req.ThreatType, req.VersionToken, req.Constraints.SupportedCompressions)
+		resp, err := api.ListUpdate(ctx, req)
 		if err != nil {
 			db.log.Printf("ListUpdate failure (%d): %v", db.updateAPIErrors+1, err)
 			db.setError(err)

webrisk_client.go

@@ -94,8 +94,9 @@ const (
 
 // Errors specific to this package.
 var (
-	errClosed = errors.New("webrisk: handler is closed")
-	errStale  = errors.New("webrisk: threat list is stale")
+	errClosed     = errors.New("webrisk: handler is closed")
+	errStale      = errors.New("webrisk: threat list is stale")
+	errMaxEntries = errors.New("webrisk: max entries must be a power of 2 between 2 ** 10 and 2 ** 20")
 )
 
 // ThreatType is an enumeration type for threats classes. Examples of threat
@@ -167,6 +168,18 @@ type Config struct {
 	// If empty, ThreatLists will be loaded instead.
 	ThreatListArg string
 
+	// MaxDiffEntries sets the maximum entries to request in a single call to ComputeThreatListDiff.
+	// This can be used in resource-constrained environments to limit the number of entries fetched
+	// at once. The default behavior (0) is to ignore this limit.
+	// If set, this should be a power of 2 between 2 ** 10 and 2 ** 20.
+	MaxDiffEntries int32
+
+	// MaxDatabaseEntries sets the maximum entries that the client will accept & store in the local
+	// database. This can be used to limit the size of the blocklist at the trade-off of decreased
+	// coverage. The default behavior (0) is to ignore this limit.
+	// If set, this should be a power of 2 between 2 ** 10 and 2 ** 20.
+	MaxDatabaseEntries int32
+
 	// ThreatLists determines which threat lists that UpdateClient should
 	// subscribe to. The threats reported by LookupURLs will only be ones that
 	// are specified by this list.
@@ -229,6 +242,19 @@ func parseThreatTypes(args string) ([]ThreatType, error) {
 	return r, nil
 }
 
+// validateMaxEntries validates a max entries argument, which must be either 0 or a power of 2
+// between 2 ** 10 and 2 ** 20.
+func validateMaxEntries(n int32) error {
+	if n == 0 {
+		return nil
+	}
+	// Bitwise check confirms a power of 2.
+	if n&(n-1) != 0 || n < 1024 || n > 1048576 {
+		return errMaxEntries
+	}
+	return nil
+}
+
 func (c Config) copy() Config {
 	c2 := c
 	c2.ThreatLists = append([]ThreatType(nil), c.ThreatLists...)
@@ -287,6 +313,11 @@ func NewUpdateClient(conf Config) (*UpdateClient, error) {
 		conf.ThreatLists = tl
 	}
 
+	// Validate max entries if args are passed.
+	if err := validateMaxEntries(conf.MaxDiffEntries); err != nil {
+		return nil, err
+	}
+
 	// Create the SafeBrowsing object.
 	if conf.api == nil {
 		var err error

webrisk_client_system_test.go

@@ -42,8 +42,7 @@ func TestNetworkAPIUpdate(t *testing.T) {
 		ThreatType: pb.ThreatType_MALWARE,
 	}
 
-	dat, err := nm.ListUpdate(context.Background(), req.ThreatType,
-		req.VersionToken, []pb.CompressionType{})
+	dat, err := nm.ListUpdate(context.Background(), req)
 	if err != nil {
 		t.Fatal(err)
 	}

webrisk_client_test.go

@@ -65,3 +65,50 @@ func TestParseThreatTypes(t *testing.T) {
 		}
 	}
 }
+
+func TestValidateMaxEntries(t *testing.T) {
+	tests := []struct {
+		n       int32
+		wantErr error
+	}{
+		{
+			n:       0,
+			wantErr: nil,
+		},
+		{
+			n:       1024,
+			wantErr: nil,
+		},
+		{
+			n:       4096,
+			wantErr: nil,
+		},
+		{
+			n:       1048576,
+			wantErr: nil,
+		},
+		{
+			n:       -1024,
+			wantErr: errMaxEntries,
+		},
+		{
+			n:       100,
+			wantErr: errMaxEntries,
+		},
+		{
+			n:       1026,
+			wantErr: errMaxEntries,
+		},
+		{
+			n:       2097152,
+			wantErr: errMaxEntries,
+		},
+	}
+
+	for _, tc := range tests {
+		gotErr := validateMaxEntries(tc.n)
+		if gotErr != tc.wantErr {
+			t.Errorf("validateMaxEntries(%d) = %v, want %v", tc.n, gotErr, tc.wantErr)
+		}
+	}
+}