add document ttl

4 files changed, 52 insertions, 9 deletions

diff --git a/api/routes.go b/api/routes.go
index ff65e89..760ee35 100644
--- a/api/routes.go
+++ b/api/routes.go
@@ -3,11 +3,11 @@ package api
 import (
 	"encoding/json"
 	"fmt"
-	"github.com/gorilla/mux"
-	"github.com/jackyzha0/ctrl-v/cache"
 	"net/http"
 	"time"
 
+	"github.com/gorilla/mux"
+	"github.com/jackyzha0/ctrl-v/cache"
 	"github.com/jackyzha0/ctrl-v/db"
 
 	log "github.com/sirupsen/logrus"
@@ -20,6 +20,7 @@ func healthCheckFunc(w http.ResponseWriter, r *http.Request) {
 func insertFunc(w http.ResponseWriter, r *http.Request) {
 	// get content
 	_ = r.ParseMultipartForm(0)
+	expiry := r.FormValue("expiry")
 	content := r.FormValue("content")
 
 	// get ip
@@ -28,7 +29,7 @@ func insertFunc(w http.ResponseWriter, r *http.Request) {
 	log.Infof("got content '%s' and ip '%s'", content, ip)
 
 	// insert content
-	err := db.New(ip, content)
+	err := db.New(ip, content, expiry)
 	if err != nil {
 		fmt.Fprintf(w, "got err: %s", err.Error())
 	}
@@ -47,10 +48,11 @@ func getHashFunc(w http.ResponseWriter, r *http.Request) {
 
 	// otherwise, return paste content and current time
 	w.Header().Set("Content-Type", "application/json")
-	pasteMap := map[string]interface{} {
+	pasteMap := map[string]interface{}{
 		"timestamp": time.Now(),
-		"content": paste.Content,
+		"content":   paste.Content,
 	}
+
 	jsonData, _ := json.Marshal(pasteMap)
 	fmt.Fprintf(w, "%+v", string(jsonData))
 }
diff --git a/db/db.go b/db/db.go
index 4a932f5..053ba87 100644
--- a/db/db.go
+++ b/db/db.go
@@ -1,7 +1,9 @@
 package db
 
 import (
+	"fmt"
 	"os"
+	"time"
 
 	"github.com/jackyzha0/ctrl-v/hashing"
 	"github.com/joho/godotenv"
@@ -23,7 +25,7 @@ func init() {
 }
 
 // creates a new paste with content and hash
-func New(ip, content string) error {
+func New(ip, content, expiry string) error {
 	// generate hash from ip
 	hash := hashing.GenerateURI(ip)
 
@@ -33,6 +35,27 @@ func New(ip, content string) error {
 		Content: content,
 	}
 
+	// check if expiry
+	if expiry != "" {
+		t, err := time.Parse(time.RFC3339, expiry)
+
+		// if time format not current
+		if err != nil {
+			return err
+		}
+
+		// time is in the past
+		if t.After(time.Now()) {
+			return fmt.Errorf("err: time %s is in the past", t.String())
+		}
+
+		new.Expiry = t
+
+	} else {
+		// 5 year expiry
+		new.Expiry = time.Now().Add(time.Hour * 43800)
+	}
+
 	// insert struct
 	log.Infof("create new paste with hash %s", hash)
 	insertErr := insert(new)
diff --git a/db/mongo.go b/db/mongo.go
index 83e71f7..4c8a739 100644
--- a/db/mongo.go
+++ b/db/mongo.go
@@ -35,6 +35,21 @@ func initSessions(user, pass, ip string) {
 		log.Fatalf("error establishing connection to mongo: %s", err.Error())
 	}
 
+	// ensure expiry check
+	sessionTTL := mgo.Index{
+		Key:         []string{"expiry"},
+		ExpireAfter: 0,
+	}
+
+	// ensure hashes are unique
+	uniqueHashes := mgo.Index{
+		Key:    []string{"hash"},
+		Unique: true,
+	}
+
+	_ = Session.DB("main").C("pastes").EnsureIndex(sessionTTL)
+	_ = Session.DB("main").C("pastes").EnsureIndex(uniqueHashes)
+
 	// Define connection to Databases
 	pastes = Session.DB("main").C("pastes")
 }
diff --git a/db/schemas.go b/db/schemas.go
index bad1e53..bdfa60c 100644
--- a/db/schemas.go
+++ b/db/schemas.go
@@ -1,12 +1,15 @@
 package db
 
-import "github.com/globalsign/mgo/bson"
+import (
+	"time"
+
+	"github.com/globalsign/mgo/bson"
+)
 
 // Paste represents a single paste
 type Paste struct {
 	ID      bson.ObjectId `bson:"_id,omitempty"`
 	Hash    string
 	Content string
+	Expiry  time.Time `bson:"expiry"`
 }
-
-// add timestamp and expiry later