add document ttl

3 files changed, 45 insertions, 4 deletions

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