From 668aad061aa3b94aa63cb8357e6c42809c3588c7 Mon Sep 17 00:00:00 2001
From: Andreas Auernhammer <github@aead.dev>
Date: Wed, 24 Jan 2024 16:53:21 +0100
Subject: [PATCH] kms: add `BackupDB` API

This commit adds the `BackupDB` client API.
It returns a snapshot of the server database
and uses compression to reduce network traffic.

Signed-off-by: Andreas Auernhammer <github@aead.dev>
---
 kms/client.go                  | 48 ++++++++++++++++++++++++++++++++++
 kms/internal/headers/header.go |  8 +++++-
 kms/request.go                 |  4 +++
 kms/response.go                | 39 +++++++++++++++++++++++++++
 4 files changed, 98 insertions(+), 1 deletion(-)

diff --git a/kms/client.go b/kms/client.go
index 3e4b96c..8cf5b05 100644
--- a/kms/client.go
+++ b/kms/client.go
@@ -6,6 +6,7 @@ package kms
 
 import (
 	"bytes"
+	"compress/gzip"
 	"context"
 	"crypto/rand"
 	"crypto/tls"
@@ -326,6 +327,53 @@ func (c *Client) RemoveNode(ctx context.Context, req *RemoveNodeRequest) error {
 	return nil
 }
 
+// BackupDB returns an io.ReadCloser containing a snapshot of the
+// current KMS server database state. The returned BackupDBResponse
+// must be closed by the caller.
+func (c *Client) BackupDB(ctx context.Context, _ *BackupDBRequest) (*BackupDBResponse, error) {
+	const (
+		Method   = http.MethodGet
+		Path     = api.PathClusterBackup
+		StatusOK = http.StatusOK
+	)
+
+	url, err := c.lb.URL(Path)
+	if err != nil {
+		return nil, err
+	}
+	r, err := http.NewRequestWithContext(ctx, Method, url, nil)
+	if err != nil {
+		return nil, err
+	}
+	r.Header.Add(headers.Accept, headers.ContentTypeAppAny)
+	r.Header.Add(headers.Accept, headers.ContentEncodingGZIP)
+
+	resp, err := c.client.Do(r)
+	if err != nil {
+		return nil, err
+	}
+	if resp.StatusCode != StatusOK {
+		return nil, readError(resp)
+	}
+
+	// Decompress the response body if the HTTP client doesn't
+	// decompress automatically.
+	body := resp.Body
+	if resp.Header.Get(headers.ContentEncoding) == headers.ContentEncodingGZIP {
+		z, err := gzip.NewReader(body)
+		if err != nil {
+			return nil, err
+		}
+		body = gzipReadCloser{
+			gzip:   z,
+			closer: body,
+		}
+	}
+	return &BackupDBResponse{
+		Body: body,
+	}, nil
+}
+
 // CreateEnclave creates a new enclave with the name req.Name.
 //
 // It returns ErrEnclaveExists if such an enclave already exists.
diff --git a/kms/internal/headers/header.go b/kms/internal/headers/header.go
index 0604c57..8b39a3f 100644
--- a/kms/internal/headers/header.go
+++ b/kms/internal/headers/header.go
@@ -9,9 +9,10 @@ package headers
 const (
 	Accept           = "Accept"            // RFC 2616
 	Authorization    = "Authorization"     // RFC 2616
+	ETag             = "ETag"              // RFC 2616
 	ContentType      = "Content-Type"      // RFC 2616
 	ContentLength    = "Content-Length"    // RFC 2616
-	ETag             = "ETag"              // RFC 2616
+	ContentEncoding  = "Content-Encoding"  // RFC 2616 and 7231
 	TransferEncoding = "Transfer-Encoding" // RFC 2616
 )
 
@@ -39,3 +40,8 @@ const (
 	ContentTypeText      = "text/plain"
 	ContentTypeHTML      = "text/html"
 )
+
+// Commonly used HTTP content encoding values.
+const (
+	ContentEncodingGZIP = "gzip"
+)
diff --git a/kms/request.go b/kms/request.go
index e467cb7..d4e29e8 100644
--- a/kms/request.go
+++ b/kms/request.go
@@ -100,6 +100,10 @@ func (r *EditClusterRequest) UnmarshalPB(v *pb.EditClusterRequest) error {
 	return nil
 }
 
+// BackupDBRequest contains options for requesting a database backup from
+// a KMS server.
+type BackupDBRequest struct{}
+
 // CreateEnclaveRequest contains options for creating enclaves.
 type CreateEnclaveRequest struct {
 	// Name is the name of the enclave to create.
diff --git a/kms/response.go b/kms/response.go
index acdf926..effd0f6 100644
--- a/kms/response.go
+++ b/kms/response.go
@@ -5,6 +5,8 @@
 package kms
 
 import (
+	"compress/gzip"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -286,6 +288,43 @@ func (s *ClusterStatusResponse) UnmarshalPB(v *pb.ClusterStatusResponse) error {
 	return nil
 }
 
+// BackupDBResponse contains the database content received from a KMS server.
+type BackupDBResponse struct {
+	Body io.ReadCloser // The database content
+}
+
+// Read reads data from the response body into b.
+func (r *BackupDBResponse) Read(b []byte) (int, error) {
+	n, err := r.Body.Read(b)
+	if errors.Is(err, io.EOF) {
+		r.Body.Close()
+	}
+	return n, err
+}
+
+// Close closes the underlying response body.
+func (r *BackupDBResponse) Close() error {
+	return r.Body.Close()
+}
+
+// gzipReadCloser wraps a gzip.Reader. It's Close method
+// closes the underlying HTTP response body and the gzip
+// reader.
+type gzipReadCloser struct {
+	gzip   *gzip.Reader
+	closer io.Closer
+}
+
+func (r gzipReadCloser) Read(b []byte) (int, error) { return r.gzip.Read(b) }
+
+func (r gzipReadCloser) Close() error {
+	err := r.closer.Close()
+	if gzipErr := r.gzip.Close(); err == nil {
+		return gzipErr
+	}
+	return err
+}
+
 // DescribeEnclaveResponse contains information about an enclave.
 type DescribeEnclaveResponse struct {
 	// Name is the name of the enclave.