Repository: google/exposure-notifications-server (go)

Issue:
Service to generate data
### TL;DR

Create a deployable service to generate TEKs. For testing purposes only.

### Design

**Proposal**

New HTTP handler in /cmd that can use be deployed and scheduled.

**Alternatives considered**

Current generation tools can't run unattended on cloud infrastructure.

/assign
/kind enhancement

Pull request:
add deployable service to generate data
/fixes #450 

Patch:
diff --git a/cmd/generate/main.go b/cmd/generate/main.go
new file mode 100644
index 0000000..0c8e31d
--- /dev/null
+++ b/cmd/generate/main.go
@@ -0,0 +1,52 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// 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.
+
+// Deployable server to generate fake exposure keys on demand.
+// THIS IS NOT INTENED TO RUN IN PRODUCTION.
+package main
+
+import (
+	"context"
+	"log"
+	"net/http"
+
+	"go.opencensus.io/plugin/ochttp"
+
+	"github.com/google/exposure-notifications-server/internal/generate"
+	"github.com/google/exposure-notifications-server/internal/logging"
+	_ "github.com/google/exposure-notifications-server/internal/observability"
+	"github.com/google/exposure-notifications-server/internal/setup"
+)
+
+func main() {
+	ctx := context.Background()
+	logger := logging.FromContext(ctx)
+
+	var config generate.Config
+	env, closer, err := setup.Setup(ctx, &config)
+	if err != nil {
+		logger.Fatalf("setup.Setup: %v", err)
+	}
+	defer closer()
+
+	handler, err := generate.NewHandler(ctx, &config, env)
+	if err != nil {
+		logger.Fatalf("unable to create generate handler: %v", err)
+	}
+	mux := http.NewServeMux()
+	mux.Handle("/", handler)
+	logger.Infof("starting generate server on :%s", config.Port)
+	instrumentedHandler := &ochttp.Handler{Handler: mux}
+	log.Fatal(http.ListenAndServe(":"+config.Port, instrumentedHandler))
+}
diff --git a/internal/export/exportfile.go b/internal/export/exportfile.go
index 8981b98..483fd2d 100644
--- a/internal/export/exportfile.go
+++ b/internal/export/exportfile.go
@@ -21,6 +21,7 @@ import (
 	"crypto/rand"
 	"crypto/sha256"
 	"fmt"
+	"io/ioutil"
 	"sort"
 
 	"github.com/google/exposure-notifications-server/internal/export/model"
@@ -32,7 +33,6 @@ import (
 )
 
 const (
-	fixedHeaderWidth     = 16
 	exportBinaryName     = "export.bin"
 	exportSignatureName  = "export.sig"
 	defaultIntervalCount = 144
@@ -40,6 +40,11 @@ const (
 	algorithm = "1.2.840.10045.4.3.2"
 )
 
+var (
+	fixedHeader      = []byte("EK Export v1    ")
+	fixedHeaderWidth = 16
+)
+
 type ExportSigners struct {
 	SignatureInfo *model.SignatureInfo
 	Signer        crypto.Signer
@@ -84,8 +89,50 @@ func MarshalExportFile(eb *model.ExportBatch, exposures []*publishmodel.Exposure
 	return buf.Bytes(), nil
 }
 
+// Unmarshal extracts the protobuf encoded exposure key present in the zip archived payload.
+func UnmarshalExportFile(zippedProtoPayload []byte) (*export.TemporaryExposureKeyExport, error) {
+	zp, err := zip.NewReader(bytes.NewReader(zippedProtoPayload), int64(len(zippedProtoPayload)))
+	if err != nil {
+		return nil, fmt.Errorf("Can't read payload: %v", err)
+	}
+
+	for _, file := range zp.File {
+		if file.Name == exportBinaryName {
+			return unmarshalContent(file)
+		}
+	}
+
+	return nil, fmt.Errorf("Payload is invalid: no %v file was found", exportBinaryName)
+}
+
+func unmarshalContent(file *zip.File) (*export.TemporaryExposureKeyExport, error) {
+	f, err := file.Open()
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+
+	content, err := ioutil.ReadAll(f)
+	if err != nil {
+		return nil, err
+	}
+
+	prefix := content[:fixedHeaderWidth]
+	if !bytes.Equal(prefix, fixedHeader) {
+		return nil, fmt.Errorf("unknown prefix: %v", string(prefix))
+	}
+
+	message := new(export.TemporaryExposureKeyExport)
+	err = proto.Unmarshal(content[fixedHeaderWidth:], message)
+	if err != nil {
+		return nil, err
+	}
+
+	return message, nil
+}
+
 func marshalContents(eb *model.ExportBatch, exposures []*publishmodel.Exposure, batchNum int32, batchSize int32, signers []ExportSigners) ([]byte, error) {
-	exportBytes := []byte("EK Export v1    ")
+	exportBytes := fixedHeader
 	if len(exportBytes) != fixedHeaderWidth {
 		return nil, fmt.Errorf("incorrect header length: %d", len(exportBytes))
 	}
diff --git a/internal/generate/config.go b/internal/generate/config.go
new file mode 100644
index 0000000..c0ff45a
--- /dev/null
+++ b/internal/generate/config.go
@@ -0,0 +1,45 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// 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.
+
+// Package generate contains HTTP handler for triggering data generation into the databae.
+package generate
+
+import (
+	"time"
+
+	"github.com/google/exposure-notifications-server/internal/database"
+	"github.com/google/exposure-notifications-server/internal/setup"
+)
+
+// Compile-time check to assert this config matches requirements.
+var _ setup.DBConfigProvider = (*Config)(nil)
+
+// Config represents the configuration and associated environment variables for
+// the publish components.
+type Config struct {
+	Port             string        `envconfig:"PORT" default:"8080"`
+	NumExposures     int           `envconfig:"NUM_EXPOSURES_GENERATED" default:"10"`
+	KeysPerExposure  int           `envconfig:"KEYS_PER_EXPOSURE" default:"14"`
+	MaxKeysOnPublish int           `envconfig:"MAX_KEYS_ON_PUBLISH" default:"15"`
+	MaxIntervalAge   time.Duration `envconfig:"MAX_INTERVAL_AGE_ON_PUBLISH" default:"360h"`
+	TruncateWindow   time.Duration `envconfig:"TRUNCATE_WINDOW" default:"1h"`
+	DefaultRegion    string        `envconfig:"DEFAULT_REGOIN" default:"US"`
+
+	Database *database.Config
+}
+
+// DB returns the configuration for the databse.
+func (c *Config) DB() *database.Config {
+	return c.Database
+}
diff --git a/internal/generate/handler.go b/internal/generate/handler.go
new file mode 100644
index 0000000..88751f8
--- /dev/null
+++ b/internal/generate/handler.go
@@ -0,0 +1,110 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// 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.
+
+package generate
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"strings"
+	"time"
+
+	"go.opencensus.io/trace"
+
+	"github.com/google/exposure-notifications-server/internal/logging"
+	"github.com/google/exposure-notifications-server/internal/publish/database"
+	"github.com/google/exposure-notifications-server/internal/publish/model"
+	"github.com/google/exposure-notifications-server/internal/serverenv"
+	"github.com/google/exposure-notifications-server/internal/util"
+)
+
+func NewHandler(ctx context.Context, config *Config, env *serverenv.ServerEnv) (http.Handler, error) {
+	logger := logging.FromContext(ctx)
+
+	if env.Database() == nil {
+		return nil, fmt.Errorf("missing database in server environment")
+	}
+
+	transformer, err := model.NewTransformer(config.MaxKeysOnPublish, config.MaxIntervalAge, config.TruncateWindow)
+	if err != nil {
+		return nil, fmt.Errorf("model.NewTransformer: %w", err)
+	}
+	logger.Infof("max keys per upload: %v", config.MaxKeysOnPublish)
+	logger.Infof("max interval start age: %v", config.MaxIntervalAge)
+	logger.Infof("truncate window: %v", config.TruncateWindow)
+
+	return &generateHandler{
+		serverenv:   env,
+		transformer: transformer,
+		config:      config,
+		database:    database.New(env.Database()),
+	}, nil
+}
+
+type generateHandler struct {
+	config      *Config
+	serverenv   *serverenv.ServerEnv
+	transformer *model.Transformer
+	database    *database.PublishDB
+}
+
+func (h *generateHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	ctx, span := trace.StartSpan(r.Context(), "(*generate.generateHandler).ServeHTTP")
+	defer span.End()
+	logger := logging.FromContext(ctx)
+
+	regionStr := h.config.DefaultRegion
+	if regionParams, ok := r.URL.Query()["region"]; ok && len(regionParams) > 0 {
+		regionStr = regionParams[0]
+	}
+	regions := strings.Split(regionStr, ",")
+	logger.Infof("Rquest to generate data for regions: %v", regions)
+
+	batchTime := time.Now()
+	for i := 0; i < h.config.NumExposures; i++ {
+		logger.Infof("Generating exposure %v of %v", i+1, h.config.NumExposures)
+
+		tr, err := util.RandomTransmissionRisk()
+		if err != nil {
+			tr = 0
+		}
+		publish := model.Publish{
+			Keys:           util.GenerateExposureKeys(h.config.KeysPerExposure, tr, false),
+			Regions:        regions,
+			AppPackageName: "generated.data",
+		}
+
+		exposures, err := h.transformer.TransformPublish(&publish, batchTime)
+		if err != nil {
+			message := fmt.Sprintf("Error transofmring generated exposures: %v", err)
+			span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: message})
+			w.WriteHeader(http.StatusInternalServerError)
+			w.Write([]byte(message))
+			return
+		}
+
+		err = h.database.InsertExposures(ctx, exposures)
+		if err != nil {
+			message := fmt.Sprintf("error writing exposure record: %v", err)
+			span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: message})
+			w.WriteHeader(http.StatusInternalServerError)
+			w.Write([]byte(message))
+			return
+		}
+		logger.Infof("Generated %v exposures", len(exposures))
+	}
+	w.WriteHeader(http.StatusOK)
+	w.Write([]byte("Generated exposure keys."))
+}
diff --git a/internal/util/generate.go b/internal/util/generate.go
index 32d9cd1..0653f2d 100644
--- a/internal/util/generate.go
+++ b/internal/util/generate.go
@@ -70,13 +70,19 @@ func RandomArrValue(arr []string) (string, error) {
 	return arr[n], nil
 }
 
-func GenerateExposureKeys(numKeys, tr int) []model.ExposureKey {
+func GenerateExposureKeys(numKeys, tr int, randomInterval bool) []model.ExposureKey {
 	// When publishing multiple keys - they'll be on different days.
-	intervalCount, err := RandomIntervalCount()
-	if err != nil {
-		log.Fatalf("problem with random interval: %v", err)
+	var err error
+	intervalCount := int32(144)
+	if randomInterval {
+		intervalCount, err = RandomIntervalCount()
+		if err != nil {
+			log.Fatalf("problem with random interval: %v", err)
+		}
 	}
-	intervalNumber := int32(time.Now().Unix()/600) - intervalCount
+	// Keys will normally align to UTC day boundries.
+	utcDay := time.Now().UTC().Truncate(24 * time.Hour)
+	intervalNumber := int32(utcDay.Unix()/600) - intervalCount
 	exposureKeys := make([]model.ExposureKey, numKeys)
 	for i := 0; i < numKeys; i++ {
 		transmissionRisk := tr
@@ -92,9 +98,11 @@ func GenerateExposureKeys(numKeys, tr int) []model.ExposureKey {
 		}
 
 		// Adjust interval math for next key.
-		intervalCount, err = RandomIntervalCount()
-		if err != nil {
-			log.Fatalf("problem with random interval: %v", err)
+		if randomInterval {
+			intervalCount, err = RandomIntervalCount()
+			if err != nil {
+				log.Fatalf("problem with random interval: %v", err)
+			}
 		}
 		intervalNumber -= intervalCount
 	}
diff --git a/tools/export-generate/main.go b/tools/export-generate/main.go
index 855639e..4d2dedb 100644
--- a/tools/export-generate/main.go
+++ b/tools/export-generate/main.go
@@ -118,7 +118,7 @@ func main() {
 		}
 
 		log.Printf("Generating %d random TEKs", *numKeys)
-		keys := util.GenerateExposureKeys(*numKeys, tr)
+		keys := util.GenerateExposureKeys(*numKeys, tr, false)
 		actualNumKeys = *numKeys
 
 		exposureKeys = make([]publishmodel.Exposure, actualNumKeys)
diff --git a/tools/exposure-client/main.go b/tools/exposure-client/main.go
index 73e5142..51c99e0 100644
--- a/tools/exposure-client/main.go
+++ b/tools/exposure-client/main.go
@@ -57,7 +57,7 @@ var (
 func main() {
 	flag.Parse()
 
-	exposureKeys := util.GenerateExposureKeys(*numKeys, *transmissionRiskFlag)
+	exposureKeys := util.GenerateExposureKeys(*numKeys, *transmissionRiskFlag, false)
 	regionIdx, err := util.RandomInt(len(defaultRegions))
 	if err != nil {
 		log.Printf("error getting random region: %v", err)


Test patch:
diff --git a/internal/export/exportfile_test.go b/internal/export/exportfile_test.go
new file mode 100644
index 0000000..4a3e46b
--- /dev/null
+++ b/internal/export/exportfile_test.go
@@ -0,0 +1,140 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// 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.
+
+package export
+
+import (
+	"crypto"
+	"io"
+	"testing"
+	"time"
+
+	"github.com/golang/protobuf/proto"
+	"github.com/google/exposure-notifications-server/internal/export/model"
+	"github.com/google/exposure-notifications-server/internal/pb/export"
+	publishmodel "github.com/google/exposure-notifications-server/internal/publish/model"
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+)
+
+func TestMarshalUnmarshalExportFile(t *testing.T) {
+	batchStartTime := time.Date(2020, 5, 1, 0, 0, 0, 0, time.UTC).Truncate(time.Second)
+	batchEndTime := batchStartTime.Add(1 * time.Hour)
+
+	batch := &model.ExportBatch{
+		BatchID:          1,
+		ConfigID:         1,
+		BucketName:       "test-bucket",
+		FilenameRoot:     "files",
+		StartTimestamp:   batchStartTime,
+		EndTimestamp:     batchEndTime,
+		Region:           "US",
+		Status:           "",
+		LeaseExpires:     time.Time{},
+		SignatureInfoIDs: []int64{1, 2},
+	}
+
+	exposures := []*publishmodel.Exposure{
+		{
+			ExposureKey:      []byte("ABC"),
+			Regions:          []string{"US"},
+			IntervalNumber:   18,
+			IntervalCount:    0,
+			CreatedAt:        batchStartTime,
+			LocalProvenance:  true,
+			TransmissionRisk: 8,
+		},
+		{
+			ExposureKey:      []byte("DEF"),
+			Regions:          []string{"CA"},
+			IntervalNumber:   118,
+			IntervalCount:    1,
+			CreatedAt:        batchEndTime,
+			LocalProvenance:  true,
+			TransmissionRisk: 1,
+		},
+	}
+
+	signatureInfo := &model.SignatureInfo{
+		SigningKey:        "/kms/project/key/1",
+		SigningKeyVersion: "1",
+		SigningKeyID:      "310",
+		EndTimestamp:      batchEndTime,
+	}
+
+	signer := &customTestSigner{}
+
+	blob, err := MarshalExportFile(batch, exposures, 1, 1, []ExportSigners{
+		{SignatureInfo: signatureInfo,
+			Signer: signer}})
+	if err != nil {
+		t.Fatalf("Can't marshal export file, %v", err)
+	}
+
+	got, err := UnmarshalExportFile(blob)
+	if err != nil {
+		t.Fatalf("Unmarshal failed: %v", err)
+	}
+
+	infos := []*export.SignatureInfo{
+		{
+			VerificationKeyVersion: proto.String("1"),
+			VerificationKeyId:      proto.String("310"),
+			SignatureAlgorithm:     proto.String("1.2.840.10045.4.3.2"),
+		},
+	}
+
+	keys := []*export.TemporaryExposureKey{
+		{
+			KeyData:                    []byte("ABC"),
+			TransmissionRiskLevel:      proto.Int32(8),
+			RollingStartIntervalNumber: proto.Int32(18),
+			RollingPeriod:              proto.Int32(0),
+		},
+		{
+			KeyData:                    []byte("DEF"),
+			TransmissionRiskLevel:      proto.Int32(1),
+			RollingStartIntervalNumber: proto.Int32(118),
+			RollingPeriod:              proto.Int32(1),
+		},
+	}
+
+	keyExport := &export.TemporaryExposureKeyExport{
+		StartTimestamp: proto.Uint64(uint64(batchStartTime.Unix())),
+		EndTimestamp:   proto.Uint64(uint64(batchEndTime.Unix())),
+		Region:         proto.String("US"),
+		BatchNum:       proto.Int32(1),
+		BatchSize:      proto.Int32(1),
+		SignatureInfos: infos,
+		Keys:           keys,
+	}
+
+	ignoredTemporaryExposureKeyExportFields := cmpopts.IgnoreUnexported(export.TemporaryExposureKeyExport{})
+	ignoredSignatureInfoFields := cmpopts.IgnoreUnexported(export.SignatureInfo{})
+	ignoredTemporaryExposureKeyFields := cmpopts.IgnoreUnexported(export.TemporaryExposureKey{})
+	diff := cmp.Diff(keyExport, got, ignoredTemporaryExposureKeyExportFields, ignoredSignatureInfoFields, ignoredTemporaryExposureKeyFields)
+	if diff != "" {
+		t.Fatalf("unmarshal mismatch (-want +got):\n%v", diff)
+	}
+}
+
+type customTestSigner struct {
+	sig []byte
+	pub crypto.PublicKey
+}
+
+func (s *customTestSigner) Public() crypto.PublicKey { return s.pub }
+func (s *customTestSigner) Sign(io.Reader, []byte, crypto.SignerOpts) ([]byte, error) {
+	return s.sig, nil
+}