Merge pull request #302 from jhaals/structs

Operate on structs

deploy/aws-lambda/deploy.sh

@@ -1,7 +1,10 @@
 #!/bin/bash
 set -ex 
-echo 'Installing Serverless'
-npm install -g serverless
+
+if ! hash serverless 2>/dev/null; then
+    echo 'Installing Serverless'
+    npm install -g serverless
+fi
 
 echo 'building yopass'
 GOOS=linux go build -o main

deploy/aws-lambda/main.go

@@ -39,7 +39,8 @@ func NewDynamo(tableName string) yopass.Database {
 }
 
 // Get item from dynamo
-func (d *Dynamo) Get(key string) (string, error) {
+func (d *Dynamo) Get(key string) (yopass.Secret, error) {
+	var s yopass.Secret
 	input := &dynamodb.GetItemInput{
 		Key: map[string]*dynamodb.AttributeValue{
 			"id": {
@@ -50,12 +51,19 @@ func (d *Dynamo) Get(key string) (string, error) {
 	}
 	result, err := d.svc.GetItem(input)
 	if err != nil {
-		return "", err
+		return s, err
 	}
 	if len(result.Item) == 0 {
-		return "", fmt.Errorf("Key not found in database")
+		return s, fmt.Errorf("Key not found in database")
 	}
-	return *result.Item["secret"].S, nil
+
+	if *result.Item["one_time"].BOOL {
+		if err := d.Delete(key); err != nil {
+			return s, err
+		}
+	}
+	s.Message = *result.Item["secret"].S
+	return s, nil
 }
 
 // Delete item
@@ -74,20 +82,23 @@ func (d *Dynamo) Delete(key string) error {
 }
 
 // Put item in Dynamo
-func (d *Dynamo) Put(key, value string, expiration int32) error {
+func (d *Dynamo) Put(key string, secret yopass.Secret) error {
 	input := &dynamodb.PutItemInput{
 		// TABLE GENERATED NAME
 		Item: map[string]*dynamodb.AttributeValue{
 			"id": {
 				S: aws.String(key),
 			},
 			"secret": {
-				S: aws.String(value),
+				S: aws.String(secret.Message),
+			},
+			"one_time": {
+				BOOL: aws.Bool(secret.OneTime),
 			},
 			"ttl": {
 				N: aws.String(
 					fmt.Sprintf(
-						"%d", time.Now().Unix()+int64(expiration))),
+						"%d", time.Now().Unix()+int64(secret.Expiration))),
 			},
 		},
 		TableName: aws.String(d.tableName),

pkg/yopass/database.go

@@ -1,8 +1,22 @@
 package yopass
 
+import "encoding/json"
+
+// Secret holds the encrypted message
+type Secret struct {
+	Expiration int32  `json:"expiration,omitempty"`
+	Message    string `json:"message"`
+	OneTime    bool   `json:"one_time,omitempty"`
+}
+
+// ToJSON converts a Secret to json
+func (s *Secret) ToJSON() ([]byte, error) {
+	return json.Marshal(&s)
+}
+
 // Database interface
 type Database interface {
-	Get(key string) (string, error)
-	Put(key, value string, expiration int32) error
+	Get(key string) (Secret, error)
+	Put(key string, secret Secret) error
 	Delete(key string) error
 }

pkg/yopass/memcached.go

@@ -1,6 +1,10 @@
 package yopass
 
-import "github.com/bradfitz/gomemcache/memcache"
+import (
+	"encoding/json"
+
+	"github.com/bradfitz/gomemcache/memcache"
+)
 
 // NewMemcached returns a new memcached database client
 func NewMemcached(server string) Database {
@@ -13,20 +17,38 @@ type Memcached struct {
 }
 
 // Get key in memcached
-func (m *Memcached) Get(key string) (string, error) {
+func (m *Memcached) Get(key string) (Secret, error) {
+	var s Secret
+
 	r, err := m.Client.Get(key)
 	if err != nil {
-		return "", err
+		return s, err
+	}
+
+	if err := json.Unmarshal(r.Value, &s); err != nil {
+		return s, err
+	}
+
+	if s.OneTime {
+		if err := m.Delete(key); err != nil {
+			return s, err
+		}
 	}
-	return string(r.Value), nil
+
+	return s, nil
 }
 
 // Put key in Memcached
-func (m *Memcached) Put(key, value string, expiration int32) error {
+func (m *Memcached) Put(key string, secret Secret) error {
+	data, err := secret.ToJSON()
+	if err != nil {
+		return err
+	}
+
 	return m.Client.Set(&memcache.Item{
 		Key:        key,
-		Value:      []byte(value),
-		Expiration: expiration})
+		Value:      data,
+		Expiration: secret.Expiration})
 }
 
 // Delete key from memcached

pkg/yopass/memcached_test.go

@@ -14,25 +14,20 @@ func TestMemcached(t *testing.T) {
 	m := NewMemcached(redisURL)
 
 	key := "f9fa5704-3ed2-4e60-b441-c426d3f9f3c1"
-	val := "my value"
+	secret := Secret{Message: "foo", OneTime: true}
 
-	err := m.Put(key, val, 123)
+	err := m.Put(key, secret)
 	if err != nil {
 		t.Fatalf("error in Put(): %v", err)
 	}
 
-	storedVal, err := m.Get(key)
+	storedSecret, err := m.Get(key)
 	if err != nil {
 		t.Fatalf("error in Get(): %v", err)
 	}
 
-	if storedVal != val {
-		t.Fatalf("expected value %s, got %s", val, storedVal)
-	}
-
-	err = m.Delete(key)
-	if err != nil {
-		t.Fatalf("error in Delete(): %v", err)
+	if storedSecret.Message != secret.Message {
+		t.Fatalf("expected value %s, got %s", secret.Message, storedSecret.Message)
 	}
 
 	_, err = m.Get(key)

pkg/yopass/redis.go

@@ -1,6 +1,7 @@
 package yopass
 
 import (
+	"encoding/json"
 	"time"
 
 	"github.com/go-redis/redis/v7"
@@ -22,20 +23,36 @@ type Redis struct {
 }
 
 // Get key from Redis
-func (r *Redis) Get(key string) (string, error) {
+func (r *Redis) Get(key string) (Secret, error) {
+	var s Secret
 	v, err := r.client.Get(key).Result()
 	if err != nil {
-		return "", err
+		return s, err
 	}
-	return v, nil
+
+	if err := json.Unmarshal([]byte(v), &s); err != nil {
+		return s, err
+	}
+
+	if s.OneTime {
+		if err := r.Delete(key); err != nil {
+			return s, err
+		}
+	}
+
+	return s, nil
 }
 
 // Put key to Redis
-func (r *Redis) Put(key, value string, expiration int32) error {
+func (r *Redis) Put(key string, secret Secret) error {
+	data, err := secret.ToJSON()
+	if err != nil {
+		return err
+	}
 	return r.client.Set(
 		key,
-		value,
-		time.Duration(expiration)*time.Second,
+		data,
+		time.Duration(secret.Expiration)*time.Second,
 	).Err()
 }
 

pkg/yopass/redis_test.go

@@ -17,9 +17,9 @@ func TestRedis(t *testing.T) {
 	}
 
 	key := "f9fa5704-3ed2-4e60-b441-c426d3f9f3c1"
-	val := "my value"
+	secret := Secret{Message: "foo", OneTime: true}
 
-	err = r.Put(key, val, 123)
+	err = r.Put(key, secret)
 	if err != nil {
 		t.Fatalf("error in Put(): %v", err)
 	}
@@ -29,13 +29,8 @@ func TestRedis(t *testing.T) {
 		t.Fatalf("error in Get(): %v", err)
 	}
 
-	if storedVal != val {
-		t.Fatalf("expected value %s, got %s", val, storedVal)
-	}
-
-	err = r.Delete(key)
-	if err != nil {
-		t.Fatalf("error in Delete(): %v", err)
+	if storedVal.Message != secret.Message {
+		t.Fatalf("expected value %s, got %s", secret.Message, storedVal.Message)
 	}
 
 	_, err = r.Get(key)

pkg/yopass/yopass.go

@@ -25,18 +25,12 @@ func New(db Database, maxLength int) Yopass {
 	}
 }
 
-type secret struct {
-	Expiration int32  `json:"expiration,omitempty"`
-	Message    string `json:"message"`
-	OneTime    bool   `json:"one_time,omitempty"`
-}
-
 // createSecret creates secret
 func (y *Yopass) createSecret(w http.ResponseWriter, request *http.Request) {
 	w.Header().Set("Access-Control-Allow-Origin", "*")
 
 	decoder := json.NewDecoder(request.Body)
-	var s secret
+	var s Secret
 	if err := decoder.Decode(&s); err != nil {
 		http.Error(w, `{"message": "Unable to parse json"}`, http.StatusBadRequest)
 		return
@@ -54,17 +48,11 @@ func (y *Yopass) createSecret(w http.ResponseWriter, request *http.Request) {
 
 	// Generate new UUID and store secret in memcache with specified expiration.
 	key := uuid.NewV4().String()
-	// Store secret json in database.
-	se, err := json.Marshal(s)
-	if err != nil {
-		http.Error(w, `{"message": "Failed to encode secret"}`, http.StatusBadRequest)
-		return
-	}
-
-	if err := y.db.Put(key, string(se), s.Expiration); err != nil {
+	if err := y.db.Put(key, s); err != nil {
 		http.Error(w, `{"message": "Failed to store secret in database"}`, http.StatusInternalServerError)
 		return
 	}
+
 	resp := map[string]string{"message": key}
 	jsonData, _ := json.Marshal(resp)
 	w.Write(jsonData)
@@ -74,31 +62,18 @@ func (y *Yopass) createSecret(w http.ResponseWriter, request *http.Request) {
 func (y *Yopass) getSecret(w http.ResponseWriter, request *http.Request) {
 	w.Header().Set("Access-Control-Allow-Origin", "*")
 
-	data, err := y.db.Get(mux.Vars(request)["key"])
+	secret, err := y.db.Get(mux.Vars(request)["key"])
 	if err != nil {
 		http.Error(w, `{"message": "Secret not found"}`, http.StatusNotFound)
 		return
 	}
 
-	var s secret
-	if err := json.Unmarshal([]byte(data), &s); err != nil {
-		http.Error(w, `{"message": "Failed to decode secret"}`, http.StatusNotFound)
-		return
-	}
-
-	if s.OneTime {
-		if err := y.db.Delete(mux.Vars(request)["key"]); err != nil {
-			http.Error(w, `{"message": "Failed to clear secret"}`, http.StatusInternalServerError)
-			return
-		}
-	}
-
-	resp, err := json.Marshal(s)
+	data, err := secret.ToJSON()
 	if err != nil {
-		http.Error(w, `{"message": "Failed to encode secret"}`, http.StatusBadRequest)
+		http.Error(w, `{"message": "Failed to encode secret"}`, http.StatusInternalServerError)
 		return
 	}
-	w.Write(resp)
+	w.Write(data)
 }
 
 // HTTPHandler containing all routes