Internal change

PiperOrigin-RevId: 504834508

Author: Web Risk Team

Dockerfile

@@ -8,9 +8,10 @@ RUN go mod download && go mod verify
 
 COPY . .
 
-# Uncomment the lines below to vet and test this container before building.
-# RUN go vet -v
-# RUN go test -v
+# Analyze our Go code and run all tests. To speed up building, such as during
+# development, consider commenting out these lines.
+RUN go vet -v
+RUN go test -v
 
 RUN CGO_ENABLED=0 go build -o /go/bin/wrserver cmd/wrserver/main.go
 
@@ -19,5 +20,5 @@ FROM gcr.io/distroless/static-debian11 as wrserver
 COPY --from=build /go/bin/wrserver /
 
 # The APIKEY Environmental Variable should be passed in at runtime. Example:
-# docker run -e APIKEY=XXXXXXXXXXXXXXXXXXXXX -p 8080:8080
-ENTRYPOINT [ "/wrserver"]
+# docker run -e APIKEY=XXXXXXXXXXXXXXXXXXXXX -p 8080:8080 <container label>
+ENTRYPOINT [ "/wrserver"]

LICENSE

@@ -199,4 +199,4 @@
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
-   limitations under the License.
+   limitations under the License.

api.go

@@ -1,4 +1,4 @@
-// Copyright 2019 Google LLC
+// Copyright 2023 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -22,24 +22,24 @@ import (
 	"net/url"
 	"strings"
 
+	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/proto"
 	pb "github.com/google/webrisk/internal/webrisk_proto"
-	"github.com/golang/protobuf/proto"
-	"github.com/golang/protobuf/jsonpb"
 )
 
 const (
-	findHashPath    = "v1/hashes:search"
-	fetchUpdatePath = "v1/threatLists:computeDiff"
-	threatTypeString = "threat_type"
-	versionTokenString = "version_token"
+	findHashPath                = "v1/hashes:search"
+	fetchUpdatePath             = "v1/threatLists:computeDiff"
+	threatTypeString            = "threat_type"
+	versionTokenString          = "version_token"
 	supportedCompressionsString = "constraints.supported_compressions"
-	hashPrefixString = "hash_prefix"
-	threatTypesString = "threat_types"
+	hashPrefixString            = "hash_prefix"
+	threatTypesString           = "threat_types"
 )
 
 // The api interface specifies wrappers around the Web Risk API.
 type api interface {
-	ListUpdate(ctx context.Context, threat_type pb.ThreatType, version_token []byte,
+	ListUpdate(ctx context.Context, threatType pb.ThreatType, versionToken []byte,
 		compressionTypes []pb.CompressionType) (*pb.ComputeThreatListDiffResponse, error)
 	HashLookup(ctx context.Context, hashPrefix []byte,
 		threatTypes []pb.ThreatType) (*pb.SearchHashesResponse, error)
@@ -67,11 +67,11 @@ func newNetAPI(root string, key string, proxy string) (*netAPI, error) {
 	httpClient := &http.Client{}
 
 	if proxy != "" {
-		proxyUrl, err := url.Parse(proxy)
+		proxyURL, err := url.Parse(proxy)
 		if err != nil {
 			return nil, err
 		}
-		httpClient = &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(proxyUrl)}}
+		httpClient = &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(proxyURL)}}
 	}
 
 	q := u.Query()
@@ -84,6 +84,9 @@ func newNetAPI(root string, key string, proxy string) (*netAPI, error) {
 // response body payload as resp.
 func (a *netAPI) doRequest(ctx context.Context, urlString string, resp proto.Message) error {
 	httpReq, err := http.NewRequest("GET", urlString, nil)
+	if err != nil {
+		return err
+	}
 	httpReq.Header.Add("Content-Type", "application/json")
 	httpReq.Header.Add("User-Agent", "Webrisk-Client/0.1.3")
 	httpReq = httpReq.WithContext(ctx)
@@ -99,7 +102,7 @@ func (a *netAPI) doRequest(ctx context.Context, urlString string, resp proto.Mes
 	if err != nil {
 		return err
 	}
-	return jsonpb.UnmarshalString(string(body), resp)
+	return protojson.Unmarshal(body, resp)
 }
 
 // ListUpdate issues a ComputeThreatListDiff API call and returns the response.

api_test.go

@@ -1,4 +1,4 @@
-// Copyright 2019 Google LLC
+// Copyright 2023 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -21,9 +21,9 @@ import (
 	"reflect"
 	"testing"
 
+	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/proto"
 	pb "github.com/google/webrisk/internal/webrisk_proto"
-	"github.com/golang/protobuf/jsonpb"
-	"github.com/golang/protobuf/proto"
 )
 
 type mockAPI struct {
@@ -51,47 +51,46 @@ func TestNetAPI(t *testing.T) {
 	var gotResp, wantResp proto.Message
 	responseMisformatter := ""
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		var p string
+		var p []byte
 		var err error
-		marshaler := jsonpb.Marshaler{}
 		for key, value := range r.URL.Query() {
-			if (key == "threat_type") {
-				if (len(value) == 0) {
+			if key == "threat_type" {
+				if len(value) == 0 {
 					t.Fatalf("missing value for key: %v", key)
 				}
 				gotReqThreatType = pb.ThreatType(pb.ThreatType_value[value[0]])
-			} else if (key == "constraints.supported_compressions") {
-				if (len(value) == 0) {
+			} else if key == "constraints.supported_compressions" {
+				if len(value) == 0 {
 					t.Fatalf("missing value for key: %v", key)
 				}
 				for _, comp := range value {
 					gotReqCompressionTypes = append(gotReqCompressionTypes,
 						pb.CompressionType(pb.CompressionType_value[comp]))
 				}
-			} else if (key == "hash_prefix") {
-				if (len(value) == 0) {
+			} else if key == "hash_prefix" {
+				if len(value) == 0 {
 					t.Fatalf("missing value for key: %v", key)
 				}
 				gotReqHashPrefix, err = base64.StdEncoding.DecodeString(value[0])
-				if (err != nil) {
+				if err != nil {
 					t.Fatalf("unexpected hash prefix decoding error for: %v", value[0])
 				}
-			} else if (key == "threat_types") {
-				if (len(value) == 0) {
+			} else if key == "threat_types" {
+				if len(value) == 0 {
 					t.Fatalf("missing value for key: %v", key)
 				}
 				for _, threat := range value {
 					gotReqThreatTypes = append(gotReqThreatTypes,
 						pb.ThreatType(pb.ThreatType_value[threat]))
 				}
-			} else if (key != "key") {
+			} else if key != "key" {
 				t.Fatalf("unexpected request param error for key: %v", key)
 			}
 		}
-		if p, err = marshaler.MarshalToString(wantResp); err != nil {
-			t.Fatalf("unexpected jsonpb MarshalToString error: %v", err)
+		if p, err = protojson.Marshal(wantResp); err != nil {
+			t.Fatalf("unexpected json MarshalToString error: %v", err)
 		}
-		if _, err := w.Write([]byte(responseMisformatter + p)); err != nil {
+		if _, err := w.Write([]byte(responseMisformatter + string(p))); err != nil {
 			t.Fatalf("unexpected ResponseWriter.Write error: %v", err)
 		}
 	}))
@@ -113,8 +112,8 @@ func TestNetAPI(t *testing.T) {
 			RawIndices: &pb.RawIndices{Indices: []int32{1, 2, 3}},
 		},
 	}
-	resp1, err := api.ListUpdate(context.Background(), wantReqThreatType, []byte {},
-	    wantReqCompressionTypes)
+	resp1, err := api.ListUpdate(context.Background(), wantReqThreatType, []byte{},
+		wantReqCompressionTypes)
 	gotResp = resp1
 	if err != nil {
 		t.Errorf("unexpected ListUpdate error: %v", err)
@@ -127,7 +126,7 @@ func TestNetAPI(t *testing.T) {
 		t.Errorf("mismatching ListUpdate requests for compression types:\ngot  %+v\nwant %+v",
 			gotReqCompressionTypes, wantReqCompressionTypes)
 	}
-	if !reflect.DeepEqual(gotResp, wantResp) {
+	if !proto.Equal(gotResp, wantResp) {
 		t.Errorf("mismatching ListUpdate responses:\ngot  %+v\nwant %+v", gotResp, wantResp)
 	}
 
@@ -151,7 +150,7 @@ func TestNetAPI(t *testing.T) {
 		t.Errorf("mismatching HashLookup requests for threat types:\ngot  %+v\nwant %+v",
 			gotReqThreatTypes, wantReqThreatTypes)
 	}
-	if !reflect.DeepEqual(gotResp, wantResp) {
+	if !proto.Equal(gotResp, wantResp) {
 		t.Errorf("mismatching HashLookup responses:\ngot  %+v\nwant %+v", gotResp, wantResp)
 	}
 

cache.go

@@ -1,4 +1,4 @@
-// Copyright 2019 Google LLC
+// Copyright 2023 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -14,12 +14,10 @@
 package webrisk
 
 import (
-	"fmt"
 	"sync"
 	"time"
 
 	pb "github.com/google/webrisk/internal/webrisk_proto"
-	pt "github.com/golang/protobuf/ptypes"
 )
 
 type cacheResult int
@@ -85,17 +83,13 @@ func (c *cache) Update(req *pb.SearchHashesRequest, resp *pb.SearchHashesRespons
 			c.pttls[fullHash] = make(map[ThreatType]time.Time)
 		}
 		for _, tt := range threat.ThreatTypes {
-			var err error
-			c.pttls[fullHash][ThreatType(tt)], err = pt.Timestamp(threat.ExpireTime)
-			if err != nil {
-				return fmt.Errorf("pt.Timestamp: %v", err)
-			}
+			c.pttls[fullHash][ThreatType(tt)] = threat.ExpireTime.AsTime()
 		}
 	}
 
 	// Insert negative TTLs for partial hashes.
 	if resp.GetNegativeExpireTime() != nil {
-		nttl, _ := pt.Timestamp(resp.GetNegativeExpireTime())
+		nttl := resp.GetNegativeExpireTime().AsTime()
 		partialHash := hashPrefix(req.HashPrefix)
 		c.nttls[partialHash] = nttl
 	}

cache_test.go

@@ -1,4 +1,4 @@
-// Copyright 2019 Google LLC
+// Copyright 2023 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -18,8 +18,9 @@ import (
 	"testing"
 	"time"
 
+	timepb "google.golang.org/protobuf/types/known/timestamppb"
+
 	pb "github.com/google/webrisk/internal/webrisk_proto"
-	pt "github.com/golang/protobuf/ptypes"
 )
 
 func TestCacheLookup(t *testing.T) {
@@ -164,8 +165,9 @@ func TestCacheLookup(t *testing.T) {
 func TestCacheUpdate(t *testing.T) {
 	now := time.Unix(1451436338, 951473000)
 	mockNow := func() time.Time { return now }
-	ts, _ := pt.TimestampProto(now.Add(1000 * time.Second))
-	tft, _ := pt.Timestamp(ts)
+
+	ts := timepb.New(now.Add(1000 * time.Second))
+	tft := ts.AsTime()
 
 	vectors := []struct {
 		req       *pb.SearchHashesRequest

cmd/wrlookup/main.go

@@ -1,4 +1,4 @@
-// Copyright 2019 Google LLC
+// Copyright 2023 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -11,17 +11,18 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
-
 // Command wrlookup is a tool for looking up URLs via the command-line.
 //
 // The tool reads one URL per line from STDIN and checks every URL against
 // the Web Risk API. The "Safe" or "Unsafe" verdict is printed to STDOUT.
 // If an error occurred, debug information may be printed to STDERR.
 //
 // To build the tool:
+//
 //	$ go get github.com/google/webrisk/cmd/wrlookup
 //
 // Example usage:
+//
 //	$ wrlookup -apikey $APIKEY
 //	https://google.com
 //	Safe URL: https://google.com
@@ -33,8 +34,9 @@ import (
 	"bufio"
 	"flag"
 	"fmt"
-	"github.com/google/webrisk"
 	"os"
+
+	"github.com/google/webrisk"
 )
 
 var (
@@ -77,7 +79,7 @@ func main() {
 		fmt.Fprintln(os.Stderr, "No -apikey specified")
 		os.Exit(codeInvalid)
 	}
-	sb, err := webrisk.NewWebriskClient(webrisk.Config{
+	sb, err := webrisk.NewUpdateClient(webrisk.Config{
 		APIKey:    *apiKeyFlag,
 		DBPath:    *databaseFlag,
 		Logger:    os.Stderr,