Adds maxDiffEntries and maxDatabaseEntires flags to wrserver and …

…wrlookup. The default value of 0 means that both will be ignored unless set. PiperOrigin-RevId: 640202575
google · Jun 4, 2024 · cf2e944 · cf2e944
1 parent 658f85c
commit cf2e944
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 22 deletions.
diff --git a/cmd/wrlookup/main.go b/cmd/wrlookup/main.go
@@ -40,11 +40,13 @@ import (
 )
 
 var (
-	apiKeyFlag      = flag.String("apikey", "", "specify your Web Risk API key")
-	databaseFlag    = flag.String("db", "", "path to the Web Risk database. By default persistent storage is disabled (not recommended).")
-	serverURLFlag   = flag.String("server", webrisk.DefaultServerURL, "Web Risk API server address.")
-	proxyFlag       = flag.String("proxy", "", "proxy to use to connect to the HTTP server")
-	threatTypesFlag = flag.String("threatTypes", "ALL", "threat types to check against")
+	apiKeyFlag             = flag.String("apikey", "", "specify your Web Risk API key")
+	databaseFlag           = flag.String("db", "", "path to the Web Risk database. By default persistent storage is disabled (not recommended).")
+	serverURLFlag          = flag.String("server", webrisk.DefaultServerURL, "Web Risk API server address.")
+	proxyFlag              = flag.String("proxy", "", "proxy to use to connect to the HTTP server")
+	threatTypesFlag        = flag.String("threatTypes", "ALL", "threat types to check against")
+	maxDiffEntriesFlag     = flag.Int("maxDiffEntries", 0, "maximum number of diff entries to return from a ComputeThreatListDiff request")
+	maxDatabaseEntriesFlag = flag.Int("maxDatabaseEntries", 0, "maximum number of database entries to be stored in the local database")
 )
 
 const usage = `wrlookup: command-line tool to lookup URLs with Web Risk.
@@ -81,12 +83,14 @@ func main() {
 		os.Exit(codeInvalid)
 	}
 	sb, err := webrisk.NewUpdateClient(webrisk.Config{
-		APIKey:        *apiKeyFlag,
-		DBPath:        *databaseFlag,
-		Logger:        os.Stderr,
-		ServerURL:     *serverURLFlag,
-		ProxyURL:      *proxyFlag,
-		ThreatListArg: *threatTypesFlag,
+		APIKey:             *apiKeyFlag,
+		DBPath:             *databaseFlag,
+		Logger:             os.Stderr,
+		ServerURL:          *serverURLFlag,
+		ProxyURL:           *proxyFlag,
+		ThreatListArg:      *threatTypesFlag,
+		MaxDiffEntries:     int32(*maxDiffEntriesFlag),
+		MaxDatabaseEntries: int32(*maxDatabaseEntriesFlag),
 	})
 	if err != nil {
 		fmt.Fprintln(os.Stderr, "Unable to initialize Web Risk client: ", err)

diff --git a/cmd/wrserver/main.go b/cmd/wrserver/main.go
@@ -223,11 +223,13 @@ const (
 )
 
 var (
-	apiKeyFlag      = flag.String("apikey", os.Getenv("APIKEY"), "specify your Web Risk API key")
-	srvAddrFlag     = flag.String("srvaddr", "0.0.0.0:8080", "TCP network address the HTTP server should use")
-	proxyFlag       = flag.String("proxy", "", "proxy to use to connect to the HTTP server")
-	databaseFlag    = flag.String("db", "", "path to the Web Risk database.")
-	threatTypesFlag = flag.String("threatTypes", "ALL", "threat types to check against")
+	apiKeyFlag             = flag.String("apikey", os.Getenv("APIKEY"), "specify your Web Risk API key")
+	srvAddrFlag            = flag.String("srvaddr", "0.0.0.0:8080", "TCP network address the HTTP server should use")
+	proxyFlag              = flag.String("proxy", "", "proxy to use to connect to the HTTP server")
+	databaseFlag           = flag.String("db", "", "path to the Web Risk database.")
+	threatTypesFlag        = flag.String("threatTypes", "ALL", "threat types to check against")
+	maxDiffEntriesFlag     = flag.Int("maxDiffEntries", 0, "maximum number of diff entries to return from a ComputeThreatListDiff request")
+	maxDatabaseEntriesFlag = flag.Int("maxDatabaseEntries", 0, "maximum number of database entries to be stored in the local database")
 )
 
 var threatTemplate = map[webrisk.ThreatType]string{
@@ -477,7 +479,7 @@ func runServer(srv *http.Server) (chan os.Signal, <-chan struct{}) {
 	// start listening for interrupts
 	exit := make(chan os.Signal, 1)
 	down := make(chan struct{})
-	
+
 	// runs shutdown and cleanup on an exit signal
 	go func() {
 		<-exit
@@ -518,11 +520,13 @@ func main() {
 		os.Exit(1)
 	}
 	conf := webrisk.Config{
-		APIKey:        *apiKeyFlag,
-		ProxyURL:      *proxyFlag,
-		DBPath:        *databaseFlag,
-		ThreatListArg: *threatTypesFlag,
-		Logger:        os.Stderr,
+		APIKey:             *apiKeyFlag,
+		ProxyURL:           *proxyFlag,
+		DBPath:             *databaseFlag,
+		ThreatListArg:      *threatTypesFlag,
+		MaxDiffEntries:     int32(*maxDiffEntriesFlag),
+		MaxDatabaseEntries: int32(*maxDatabaseEntriesFlag),
+		Logger:             os.Stderr,
 	}
 	wr, err := webrisk.NewUpdateClient(conf)
 	if err != nil {