Repository: google/exposure-notifications-server
Pull request #761 (2020-07-21 01:39:54)

Internal structure for revision token
* Add proto definition for revision token
* Add ability to encrypt/decrype the revision token data

Towards #734

/cc @tlepoint 

google-oss-robot: @mikehelmick: GitHub didn't allow me to request PR reviews from the following users: tlepoint.

Note that only [google members](https://github.com/orgs/google/people) and repo collaborators can review this PR, and authors cannot review their own PRs.

<details>

In response to [this](https://github.com/google/exposure-notifications-server/pull/761):

>* Add proto definition for revision token
>* Add ability to encrypt/decrype the revision token data
>
>Towards #734
>
>/cc @tlepoint 


Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository.
</details>
google-oss-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/google/exposure-notifications-server/pull/761#" title="Author self-approved">mikehelmick</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=google%2Fexposure-notifications-server).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/google/exposure-notifications-server/blob/main/OWNERS)~~ [mikehelmick]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
google-oss-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/google/exposure-notifications-server/pull/761#" title="Author self-approved">mikehelmick</a>*, *<a href="https://github.com/google/exposure-notifications-server/pull/761#pullrequestreview-452068843" title="LGTM">sethvargo</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=google%2Fexposure-notifications-server).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/google/exposure-notifications-server/blob/main/OWNERS)~~ [mikehelmick,sethvargo]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->

Patch:
diff --git a/internal/pb/revision.pb.go b/internal/pb/revision.pb.go
new file mode 100644
index 0000000..3cae079
--- /dev/null
+++ b/internal/pb/revision.pb.go
@@ -0,0 +1,325 @@
+// 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.
+
+// Code generated by protoc-gen-go. DO NOT EDIT.
+// versions:
+// 	protoc-gen-go v1.22.0
+// 	protoc        v3.11.4
+// source: internal/pb/revision.proto
+
+package pb
+
+import (
+	reflect "reflect"
+	sync "sync"
+
+	proto "github.com/golang/protobuf/proto"
+	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
+	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
+)
+
+const (
+	// Verify that this generated code is sufficiently up-to-date.
+	_ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion)
+	// Verify that runtime/protoimpl is sufficiently up-to-date.
+	_ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20)
+)
+
+// This is a compile-time assertion that a sufficiently up-to-date version
+// of the legacy proto package is being used.
+const _ = proto.ProtoPackageIsVersion4
+
+type RevisionToken struct {
+	state         protoimpl.MessageState
+	sizeCache     protoimpl.SizeCache
+	unknownFields protoimpl.UnknownFields
+
+	Kid  string `protobuf:"bytes,1,opt,name=kid,proto3" json:"kid,omitempty"`
+	Data []byte `protobuf:"bytes,2,opt,name=data,proto3" json:"data,omitempty"`
+}
+
+func (x *RevisionToken) Reset() {
+	*x = RevisionToken{}
+	if protoimpl.UnsafeEnabled {
+		mi := &file_internal_pb_revision_proto_msgTypes[0]
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		ms.StoreMessageInfo(mi)
+	}
+}
+
+func (x *RevisionToken) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*RevisionToken) ProtoMessage() {}
+
+func (x *RevisionToken) ProtoReflect() protoreflect.Message {
+	mi := &file_internal_pb_revision_proto_msgTypes[0]
+	if protoimpl.UnsafeEnabled && x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use RevisionToken.ProtoReflect.Descriptor instead.
+func (*RevisionToken) Descriptor() ([]byte, []int) {
+	return file_internal_pb_revision_proto_rawDescGZIP(), []int{0}
+}
+
+func (x *RevisionToken) GetKid() string {
+	if x != nil {
+		return x.Kid
+	}
+	return ""
+}
+
+func (x *RevisionToken) GetData() []byte {
+	if x != nil {
+		return x.Data
+	}
+	return nil
+}
+
+type RevisionTokenData struct {
+	state         protoimpl.MessageState
+	sizeCache     protoimpl.SizeCache
+	unknownFields protoimpl.UnknownFields
+
+	RevisableKeys []*RevisableKey `protobuf:"bytes,1,rep,name=revisableKeys,proto3" json:"revisableKeys,omitempty"`
+}
+
+func (x *RevisionTokenData) Reset() {
+	*x = RevisionTokenData{}
+	if protoimpl.UnsafeEnabled {
+		mi := &file_internal_pb_revision_proto_msgTypes[1]
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		ms.StoreMessageInfo(mi)
+	}
+}
+
+func (x *RevisionTokenData) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*RevisionTokenData) ProtoMessage() {}
+
+func (x *RevisionTokenData) ProtoReflect() protoreflect.Message {
+	mi := &file_internal_pb_revision_proto_msgTypes[1]
+	if protoimpl.UnsafeEnabled && x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use RevisionTokenData.ProtoReflect.Descriptor instead.
+func (*RevisionTokenData) Descriptor() ([]byte, []int) {
+	return file_internal_pb_revision_proto_rawDescGZIP(), []int{1}
+}
+
+func (x *RevisionTokenData) GetRevisableKeys() []*RevisableKey {
+	if x != nil {
+		return x.RevisableKeys
+	}
+	return nil
+}
+
+type RevisableKey struct {
+	state         protoimpl.MessageState
+	sizeCache     protoimpl.SizeCache
+	unknownFields protoimpl.UnknownFields
+
+	TemporaryExposureKey []byte `protobuf:"bytes,1,opt,name=temporaryExposureKey,proto3" json:"temporaryExposureKey,omitempty"`
+	IntervalNumber       int32  `protobuf:"varint,2,opt,name=intervalNumber,proto3" json:"intervalNumber,omitempty"`
+	IntervalCount        int32  `protobuf:"varint,3,opt,name=intervalCount,proto3" json:"intervalCount,omitempty"`
+}
+
+func (x *RevisableKey) Reset() {
+	*x = RevisableKey{}
+	if protoimpl.UnsafeEnabled {
+		mi := &file_internal_pb_revision_proto_msgTypes[2]
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		ms.StoreMessageInfo(mi)
+	}
+}
+
+func (x *RevisableKey) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*RevisableKey) ProtoMessage() {}
+
+func (x *RevisableKey) ProtoReflect() protoreflect.Message {
+	mi := &file_internal_pb_revision_proto_msgTypes[2]
+	if protoimpl.UnsafeEnabled && x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use RevisableKey.ProtoReflect.Descriptor instead.
+func (*RevisableKey) Descriptor() ([]byte, []int) {
+	return file_internal_pb_revision_proto_rawDescGZIP(), []int{2}
+}
+
+func (x *RevisableKey) GetTemporaryExposureKey() []byte {
+	if x != nil {
+		return x.TemporaryExposureKey
+	}
+	return nil
+}
+
+func (x *RevisableKey) GetIntervalNumber() int32 {
+	if x != nil {
+		return x.IntervalNumber
+	}
+	return 0
+}
+
+func (x *RevisableKey) GetIntervalCount() int32 {
+	if x != nil {
+		return x.IntervalCount
+	}
+	return 0
+}
+
+var File_internal_pb_revision_proto protoreflect.FileDescriptor
+
+var file_internal_pb_revision_proto_rawDesc = []byte{
+	0x0a, 0x1a, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x6e, 0x61, 0x6c, 0x2f, 0x70, 0x62, 0x2f, 0x72, 0x65,
+	0x76, 0x69, 0x73, 0x69, 0x6f, 0x6e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x22, 0x35, 0x0a, 0x0d,
+	0x52, 0x65, 0x76, 0x69, 0x73, 0x69, 0x6f, 0x6e, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x12, 0x10, 0x0a,
+	0x03, 0x6b, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6b, 0x69, 0x64, 0x12,
+	0x12, 0x0a, 0x04, 0x64, 0x61, 0x74, 0x61, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x04, 0x64,
+	0x61, 0x74, 0x61, 0x22, 0x48, 0x0a, 0x11, 0x52, 0x65, 0x76, 0x69, 0x73, 0x69, 0x6f, 0x6e, 0x54,
+	0x6f, 0x6b, 0x65, 0x6e, 0x44, 0x61, 0x74, 0x61, 0x12, 0x33, 0x0a, 0x0d, 0x72, 0x65, 0x76, 0x69,
+	0x73, 0x61, 0x62, 0x6c, 0x65, 0x4b, 0x65, 0x79, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32,
+	0x0d, 0x2e, 0x52, 0x65, 0x76, 0x69, 0x73, 0x61, 0x62, 0x6c, 0x65, 0x4b, 0x65, 0x79, 0x52, 0x0d,
+	0x72, 0x65, 0x76, 0x69, 0x73, 0x61, 0x62, 0x6c, 0x65, 0x4b, 0x65, 0x79, 0x73, 0x22, 0x90, 0x01,
+	0x0a, 0x0c, 0x52, 0x65, 0x76, 0x69, 0x73, 0x61, 0x62, 0x6c, 0x65, 0x4b, 0x65, 0x79, 0x12, 0x32,
+	0x0a, 0x14, 0x74, 0x65, 0x6d, 0x70, 0x6f, 0x72, 0x61, 0x72, 0x79, 0x45, 0x78, 0x70, 0x6f, 0x73,
+	0x75, 0x72, 0x65, 0x4b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x14, 0x74, 0x65,
+	0x6d, 0x70, 0x6f, 0x72, 0x61, 0x72, 0x79, 0x45, 0x78, 0x70, 0x6f, 0x73, 0x75, 0x72, 0x65, 0x4b,
+	0x65, 0x79, 0x12, 0x26, 0x0a, 0x0e, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x76, 0x61, 0x6c, 0x4e, 0x75,
+	0x6d, 0x62, 0x65, 0x72, 0x18, 0x02, 0x20, 0x01, 0x28, 0x05, 0x52, 0x0e, 0x69, 0x6e, 0x74, 0x65,
+	0x72, 0x76, 0x61, 0x6c, 0x4e, 0x75, 0x6d, 0x62, 0x65, 0x72, 0x12, 0x24, 0x0a, 0x0d, 0x69, 0x6e,
+	0x74, 0x65, 0x72, 0x76, 0x61, 0x6c, 0x43, 0x6f, 0x75, 0x6e, 0x74, 0x18, 0x03, 0x20, 0x01, 0x28,
+	0x05, 0x52, 0x0d, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x76, 0x61, 0x6c, 0x43, 0x6f, 0x75, 0x6e, 0x74,
+	0x42, 0x40, 0x5a, 0x3e, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x67,
+	0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x65, 0x78, 0x70, 0x6f, 0x73, 0x75, 0x72, 0x65, 0x2d, 0x6e,
+	0x6f, 0x74, 0x69, 0x66, 0x69, 0x63, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x2d, 0x73, 0x65, 0x72,
+	0x76, 0x65, 0x72, 0x2f, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x6e, 0x61, 0x6c, 0x2f, 0x70, 0x62, 0x3b,
+	0x70, 0x62, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
+}
+
+var (
+	file_internal_pb_revision_proto_rawDescOnce sync.Once
+	file_internal_pb_revision_proto_rawDescData = file_internal_pb_revision_proto_rawDesc
+)
+
+func file_internal_pb_revision_proto_rawDescGZIP() []byte {
+	file_internal_pb_revision_proto_rawDescOnce.Do(func() {
+		file_internal_pb_revision_proto_rawDescData = protoimpl.X.CompressGZIP(file_internal_pb_revision_proto_rawDescData)
+	})
+	return file_internal_pb_revision_proto_rawDescData
+}
+
+var file_internal_pb_revision_proto_msgTypes = make([]protoimpl.MessageInfo, 3)
+var file_internal_pb_revision_proto_goTypes = []interface{}{
+	(*RevisionToken)(nil),     // 0: RevisionToken
+	(*RevisionTokenData)(nil), // 1: RevisionTokenData
+	(*RevisableKey)(nil),      // 2: RevisableKey
+}
+var file_internal_pb_revision_proto_depIdxs = []int32{
+	2, // 0: RevisionTokenData.revisableKeys:type_name -> RevisableKey
+	1, // [1:1] is the sub-list for method output_type
+	1, // [1:1] is the sub-list for method input_type
+	1, // [1:1] is the sub-list for extension type_name
+	1, // [1:1] is the sub-list for extension extendee
+	0, // [0:1] is the sub-list for field type_name
+}
+
+func init() { file_internal_pb_revision_proto_init() }
+func file_internal_pb_revision_proto_init() {
+	if File_internal_pb_revision_proto != nil {
+		return
+	}
+	if !protoimpl.UnsafeEnabled {
+		file_internal_pb_revision_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} {
+			switch v := v.(*RevisionToken); i {
+			case 0:
+				return &v.state
+			case 1:
+				return &v.sizeCache
+			case 2:
+				return &v.unknownFields
+			default:
+				return nil
+			}
+		}
+		file_internal_pb_revision_proto_msgTypes[1].Exporter = func(v interface{}, i int) interface{} {
+			switch v := v.(*RevisionTokenData); i {
+			case 0:
+				return &v.state
+			case 1:
+				return &v.sizeCache
+			case 2:
+				return &v.unknownFields
+			default:
+				return nil
+			}
+		}
+		file_internal_pb_revision_proto_msgTypes[2].Exporter = func(v interface{}, i int) interface{} {
+			switch v := v.(*RevisableKey); i {
+			case 0:
+				return &v.state
+			case 1:
+				return &v.sizeCache
+			case 2:
+				return &v.unknownFields
+			default:
+				return nil
+			}
+		}
+	}
+	type x struct{}
+	out := protoimpl.TypeBuilder{
+		File: protoimpl.DescBuilder{
+			GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
+			RawDescriptor: file_internal_pb_revision_proto_rawDesc,
+			NumEnums:      0,
+			NumMessages:   3,
+			NumExtensions: 0,
+			NumServices:   0,
+		},
+		GoTypes:           file_internal_pb_revision_proto_goTypes,
+		DependencyIndexes: file_internal_pb_revision_proto_depIdxs,
+		MessageInfos:      file_internal_pb_revision_proto_msgTypes,
+	}.Build()
+	File_internal_pb_revision_proto = out.File
+	file_internal_pb_revision_proto_rawDesc = nil
+	file_internal_pb_revision_proto_goTypes = nil
+	file_internal_pb_revision_proto_depIdxs = nil
+}
diff --git a/internal/pb/revision.proto b/internal/pb/revision.proto
new file mode 100644
index 0000000..478f6a8
--- /dev/null
+++ b/internal/pb/revision.proto
@@ -0,0 +1,32 @@
+// 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.
+
+syntax = "proto3";
+
+option go_package = "github.com/google/exposure-notifications-server/internal/pb;pb";
+
+message RevisionToken {
+    string kid = 1;
+    bytes data = 2;
+}
+
+message RevisionTokenData {
+    repeated RevisableKey revisableKeys = 1;
+}
+
+message RevisableKey {
+    bytes temporaryExposureKey = 1;
+    int32 intervalNumber = 2;
+    int32 intervalCount = 3;
+}
\ No newline at end of file
diff --git a/internal/revision/revision.go b/internal/revision/revision.go
new file mode 100644
index 0000000..4bc7782
--- /dev/null
+++ b/internal/revision/revision.go
@@ -0,0 +1,137 @@
+// 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 revision defines the internal structure of the revision token
+// and utilities for marshal/unmarshal which also encrypts/decrypts the payload.
+package revision
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
+	"fmt"
+	"io"
+	"sort"
+
+	"github.com/google/exposure-notifications-server/internal/pb"
+	"github.com/google/exposure-notifications-server/internal/publish/model"
+	"google.golang.org/protobuf/proto"
+)
+
+func buildTokenBufer(eKeys []*model.Exposure) *pb.RevisionTokenData {
+	// sort the keys.
+	sort.Slice(eKeys, func(i, j int) bool {
+		return eKeys[i].ExposureKeyBase64() < eKeys[j].ExposureKeyBase64()
+	})
+	// Build the protocol buffer version of the revision token data.
+	tokenData := pb.RevisionTokenData{
+		RevisableKeys: make([]*pb.RevisableKey, 0, len(eKeys)),
+	}
+	for _, k := range eKeys {
+		pbKey := pb.RevisableKey{
+			TemporaryExposureKey: make([]byte, len(k.ExposureKey)),
+			IntervalNumber:       k.IntervalNumber,
+			IntervalCount:        k.IntervalCount,
+		}
+		copy(pbKey.TemporaryExposureKey, k.ExposureKey)
+		tokenData.RevisableKeys = append(tokenData.RevisableKeys, &pbKey)
+	}
+	return &tokenData
+}
+
+// MakeRevisionToken turns the TEK data from a given publish request
+// into an encrypted protocol buffer revision token.
+func MakeRevisionToken(eKeys []*model.Exposure, aad []byte, kid string, cryptoKey []byte) ([]byte, error) {
+	if len(eKeys) == 0 {
+		return nil, fmt.Errorf("no keys to build token for")
+	}
+
+	tokenData := buildTokenBufer(eKeys)
+	plaintext, err := proto.Marshal(tokenData)
+	if err != nil {
+		return nil, fmt.Errorf("unable to masrhal token data: %w", err)
+	}
+
+	// encrypt the serialized proto.
+	block, err := aes.NewCipher(cryptoKey)
+	if err != nil {
+		return nil, fmt.Errorf("bad cipher block: %w", err)
+	}
+	aesgcm, err := cipher.NewGCM(block)
+	if err != nil {
+		return nil, fmt.Errorf("failed to wrap cipher block: %w", err)
+	}
+	nonce := make([]byte, aesgcm.NonceSize())
+	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
+		return nil, fmt.Errorf("failed to generate nonce: %w", err)
+	}
+	ciphertext := aesgcm.Seal(nonce, nonce, plaintext, aad)
+
+	// Build the revision token.
+	token := pb.RevisionToken{
+		Kid:  kid,
+		Data: ciphertext,
+	}
+	tokenBytes, err := proto.Marshal(&token)
+	if err != nil {
+		return nil, fmt.Errorf("faield to marshal token: %w", err)
+	}
+
+	return tokenBytes, nil
+}
+
+// UnmarshalRevisionToken unmarshals a revision token, decrypts the payload,
+// and returns the TEK data that was contained in the token if valid.
+func UnmarshalRevisionToken(tokenBytes []byte, aad []byte, cryptoKeys map[string][]byte) (*pb.RevisionTokenData, error) {
+	var revisionToken pb.RevisionToken
+	if err := proto.Unmarshal(tokenBytes, &revisionToken); err != nil {
+		return nil, fmt.Errorf("unable to unmarshal proto envelope: %w", err)
+	}
+	data := revisionToken.Data
+
+	dek, ok := cryptoKeys[revisionToken.Kid]
+	if !ok {
+		return nil, fmt.Errorf("token has invalid key id: %v", revisionToken.Kid)
+	}
+
+	// Decrypt the data block.
+	block, err := aes.NewCipher(dek)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create cipher from dek: %w", err)
+	}
+
+	aesgcm, err := cipher.NewGCM(block)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create gcm from dek: %w", err)
+	}
+
+	size := aesgcm.NonceSize()
+	if len(data) < size {
+		return nil, fmt.Errorf("malformed ciphertext")
+	}
+	nonce, ciphertext := data[:size], data[size:]
+
+	plaintext, err := aesgcm.Open(nil, nonce, ciphertext, aad)
+	if err != nil {
+		return nil, fmt.Errorf("failed to decrypt ciphertext with dek: %w", err)
+	}
+
+	// The plaintext is a pb.RevisionTokenData
+	var tokenData pb.RevisionTokenData
+	if err := proto.Unmarshal(plaintext, &tokenData); err != nil {
+		return nil, fmt.Errorf("faield to unmarshal token data: %w", err)
+	}
+
+	return &tokenData, nil
+}


Test patch:
diff --git a/internal/revision/revision_test.go b/internal/revision/revision_test.go
new file mode 100644
index 0000000..1d5e8f6
--- /dev/null
+++ b/internal/revision/revision_test.go
@@ -0,0 +1,147 @@
+// 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 revision
+
+import (
+	"crypto/rand"
+	"io"
+	"testing"
+
+	"github.com/google/exposure-notifications-server/internal/pb"
+	"github.com/google/exposure-notifications-server/internal/publish/model"
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+)
+
+func TestBuildTokenBuffer(t *testing.T) {
+	cases := []struct {
+		name    string
+		publish []*model.Exposure
+		want    *pb.RevisionTokenData
+	}{
+		{
+			name: "one",
+			publish: []*model.Exposure{
+				{
+					ExposureKey:    []byte{1, 2, 3, 4},
+					IntervalNumber: 200000,
+					IntervalCount:  144,
+				},
+			},
+			want: &pb.RevisionTokenData{
+				RevisableKeys: []*pb.RevisableKey{
+					{
+						TemporaryExposureKey: []byte{1, 2, 3, 4},
+						IntervalNumber:       200000,
+						IntervalCount:        144,
+					},
+				},
+			},
+		},
+		{
+			name: "wrong_order",
+			publish: []*model.Exposure{
+				{
+					ExposureKey:    []byte{4, 3, 2, 1},
+					IntervalNumber: 200144,
+					IntervalCount:  144,
+				},
+				{
+					ExposureKey:    []byte{1, 2, 3, 4},
+					IntervalNumber: 200000,
+					IntervalCount:  144,
+				},
+			},
+			want: &pb.RevisionTokenData{
+				RevisableKeys: []*pb.RevisableKey{
+					{
+						TemporaryExposureKey: []byte{1, 2, 3, 4},
+						IntervalNumber:       200000,
+						IntervalCount:        144,
+					},
+					{
+						TemporaryExposureKey: []byte{4, 3, 2, 1},
+						IntervalNumber:       200144,
+						IntervalCount:        144,
+					},
+				},
+			},
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			got := buildTokenBufer(tc.publish)
+			if diff := cmp.Diff(tc.want, got, cmpopts.IgnoreUnexported(pb.RevisionTokenData{}), cmpopts.IgnoreUnexported(pb.RevisableKey{})); diff != "" {
+				t.Fatalf("mismatch (-want, +got):\n%s", diff)
+			}
+		})
+	}
+}
+
+func TestEncryptDecrypt(t *testing.T) {
+	key := make([]byte, 32)
+	if _, err := io.ReadFull(rand.Reader, key); err != nil {
+		t.Fatalf("unable to generate AES key: %v", err)
+	}
+
+	// Build data to serialize and encrypt.
+	source := []*model.Exposure{
+		{
+			ExposureKey:    []byte{1, 2, 3, 4},
+			IntervalNumber: 254321,
+			IntervalCount:  144,
+		},
+		{
+			ExposureKey:    []byte{5, 6, 7, 8},
+			IntervalNumber: 254465,
+			IntervalCount:  144,
+		},
+	}
+	aad := []byte{16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1}
+	kid := "123"
+	// Expected result in the end
+	want := &pb.RevisionTokenData{
+		RevisableKeys: []*pb.RevisableKey{
+			{
+				TemporaryExposureKey: []byte{1, 2, 3, 4},
+				IntervalNumber:       254321,
+				IntervalCount:        144,
+			},
+			{
+				TemporaryExposureKey: []byte{5, 6, 7, 8},
+				IntervalNumber:       254465,
+				IntervalCount:        144,
+			},
+		},
+	}
+
+	encrypted, err := MakeRevisionToken(source, aad, kid, key)
+	if err != nil {
+		t.Fatalf("error encrypting and serializing data: %v", err)
+	}
+
+	validKeys := map[string][]byte{
+		kid: key,
+	}
+	got, err := UnmarshalRevisionToken(encrypted, aad, validKeys)
+	if err != nil {
+		t.Fatalf("error decrypting token: %v", err)
+	}
+
+	if diff := cmp.Diff(want, got, cmpopts.IgnoreUnexported(pb.RevisionTokenData{}), cmpopts.IgnoreUnexported(pb.RevisableKey{})); diff != "" {
+		t.Fatalf("mismatch (-want, +got):\n%s", diff)
+	}
+}