3 files changed, 117 insertions, 27 deletions

diff --git a/src/http.go b/src/http.go
index 81d55b7..1e60046 100644
--- a/src/http.go
+++ b/src/http.go
@@ -36,6 +36,7 @@ func setupHTTPServer() (http.Server) {
 	r.HandleFunc("/api/startContainer", startContainerHandler).Methods("POST")
 	r.HandleFunc("/api/stopContainer", stopContainerHandler).Methods("POST")
 	r.HandleFunc("/api/getAccess", getAccessHandler).Methods("GET")
+	r.HandleFunc("/api/getTimeLimit", getTimeLimitHandler).Methods("GET")
 
 	return http.Server{
 		Addr: fmt.Sprintf("0.0.0.0:%d", *port),
@@ -113,6 +114,9 @@ func loginPostHandler(w http.ResponseWriter, r *http.Request) {
 				Expires: time.Now().Add(time.Hour * 24),
 			})
 			validRedirect = true
+
+			// register our login time for the limiter
+			registerLoginForLimiter()
 		}
 	}
 
@@ -211,30 +215,36 @@ func submitFlagHandler(w http.ResponseWriter, r *http.Request) {
 	} else {
 		// valid session token found, now search for the requested challenge
 
+		errorString := ""
 		foundChallenge := false
 		correctFlag := false
 
-		// try to find our challenge
-		for index, challenge := range challenges {
-			if challenge.Name == challengeName {
-				// found challenge, check flags
-				foundChallenge = true
-
-				if challenge.Flag == flag {
-					// our user found the flag \o/
-					challenges[index].FoundFlag = true
-					correctFlag = true
-				} else {
-					// ow, bummer :(
-					challenge.FlagTries++
+		// check if we are in the desired timeframe
+		if shouldLimit() {
+			// We are not.
+			errorString = "Time's up."
+		} else {
+			// We can check that flag. Try to find our challenge
+			for index, challenge := range challenges {
+				if challenge.Name == challengeName {
+					// found challenge, check flags
+					foundChallenge = true
+
+					if challenge.Flag == flag {
+						// our user found the flag \o/
+						challenges[index].FoundFlag = true
+						correctFlag = true
+					} else {
+						// ow, bummer :(
+						challenge.FlagTries++
+					}
+					break
 				}
-				break
 			}
 		}
 
 		// if we didn't find the challenge, write an error message
-		errorString := ""
-		if !foundChallenge {
+		if !foundChallenge && errorString != "" {
 			errorString = "no such challenge"
 		}
 
@@ -261,18 +271,26 @@ func startContainerHandler(w http.ResponseWriter, r *http.Request) {
 		http.Redirect(w, r, "/login", http.StatusTemporaryRedirect)
 	} else {
 		// valid session token found, now search for the requested challenge
-		for _, challenge := range challenges {
-			if challenge.Name == challengeName {
-				// found challenge, start container
-
-				cc, err := startChallengeContainer(challenge)
-				if err != nil {
-					log.Println(err.Error())
-					errorString = "Server error."
-				} else {
-					addressString = cc.IP
+
+		// check if we are in the desired timeframe
+		if shouldLimit() {
+			// woops! Limit starting the container.
+			errorString = "Time's up."
+		} else {
+			// we don't need to limit - start the container
+			for _, challenge := range challenges {
+				if challenge.Name == challengeName {
+					// found challenge, start container
+
+					cc, err := startChallengeContainer(challenge)
+					if err != nil {
+						log.Println(err.Error())
+						errorString = "Server error."
+					} else {
+						addressString = cc.IP
+					}
+					break
 				}
-				break
 			}
 		}
 
@@ -332,3 +350,14 @@ func getAccessHandler(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 }
+
+// Returns the configuration for the VPN
+func getTimeLimitHandler(w http.ResponseWriter, r *http.Request) {
+	// We don't need to verify session cookies.
+
+	jsonAnswer, _ := json.Marshal(map[string]string{
+		"endTimestamp": fmt.Sprintf("%d", *endTimestamp),
+		"endAfter": fmt.Sprintf("%d", *endAfter),
+	})
+	w.Write([]byte(jsonAnswer))
+}
diff --git a/src/limit.go b/src/limit.go
new file mode 100644
index 0000000..c4befac
--- /dev/null
+++ b/src/limit.go
@@ -0,0 +1,56 @@
+package main
+
+import (
+	"flag"
+	"time"
+)
+
+var (
+	endTimestamp *int64
+	endAfter *int64
+	initialLoginTime int64
+)
+
+func registerLimitFlags() {
+	endTimestamp = flag.Int64("endTimestamp", 0, "Date/Time after which flags are not accepted anymore")
+	endAfter = flag.Int64("endAfter", 0, "Seconds after the first login, after which flags are not accepted anymore")
+}
+
+func startLimitTimer() {
+	now := time.Now().Unix()
+
+	// check if endTimestamp is set
+	if *endTimestamp > 0 {
+		// Start the endTimestampTimer
+		endTimestampTimer := time.NewTimer(time.Duration(*endTimestamp - now))
+		go func() {
+			<-endTimestampTimer.C
+			// Stop all challenge containers if timer hit
+			stopAllChallengeContainers()
+		}()
+	}
+
+	// check if endAfter is set
+	if *endAfter > 0 {
+		// Start the endAfterTimer
+		endAfterTimer := time.NewTimer(time.Duration(*endAfter - now))
+		go func() {
+			<-endAfterTimer.C
+			// Stop all challenge containers if timer hit
+			stopAllChallengeContainers()
+		}()
+	}
+}
+
+// Called on every login to set the "initialLoginTime" if it's not already set
+func registerLoginForLimiter() {
+	if initialLoginTime == 0 {
+		initialLoginTime = time.Now().Unix()
+	}
+}
+
+// Called before starting containers or entering flags
+func shouldLimit() (bool) {
+	now := time.Now().Unix()
+	return (*endTimestamp > 0 && *endTimestamp < now) || (*endAfter > 0 && *endAfter < now - initialLoginTime)
+}
diff --git a/src/main.go b/src/main.go
index 91e453e..39dd213 100644
--- a/src/main.go
+++ b/src/main.go
@@ -18,6 +18,7 @@ func main() {
 	registerCredentialsFlags()
 	registerSeedFlags()
 	registerAccessFlags()
+	registerLimitFlags()
 	flag.Parse()
 
 	// Read challenges from file
@@ -38,6 +39,10 @@ func main() {
 		log.Fatalln(startVPNError.Error())
 	}
 
+	// Launch Limiter
+	log.Printf("Starting limiter (end %d, timespan %d)", *endTimestamp, *endAfter)
+	startLimitTimer()
+
 	// Set up HTTP server
 	log.Printf("Running HTTP server on port %d", *port)
 	httpServer := setupHTTPServer()