From 8e17fe45852a8c7cb19c9bd5aee749d52a24f6ff Mon Sep 17 00:00:00 2001
From: Spotlight <spotlight@joscomputing.space>
Date: Sat, 5 Sep 2020 10:34:24 -0500
Subject: [PATCH] Add proper authentication

---
 ias.go   |  2 +-
 main.go  |  1 +
 route.go | 76 +++++++++++++++++++++++++++++++++++++++++++++++++++++---
 3 files changed, 75 insertions(+), 4 deletions(-)

diff --git a/ias.go b/ias.go
index fd82d4b..aeaa3a6 100644
--- a/ias.go
+++ b/ias.go
@@ -34,7 +34,7 @@ var registerUser *sql.Stmt
 
 func iasInitialize() {
 	var err error
-	registerUser, err = db.Prepare(`INSERT INTO wiisoap.userbase (DeviceId, DeviceToken, AccountId, Region, Country, Language, SerialNo, DeviceCode)  VALUES (?, ?, ?, ?, ?, ?, ?, ?)`)
+	registerUser, err = db.Prepare(`INSERT INTO userbase (DeviceId, DeviceToken, AccountId, Region, Country, Language, SerialNo, DeviceCode)  VALUES (?, ?, ?, ?, ?, ?, ?, ?)`)
 	if err != nil {
 		log.Fatalf("ias initialize: error preparing statement: %v\n", err)
 	}
diff --git a/main.go b/main.go
index 7e988d6..04b59e6 100644
--- a/main.go
+++ b/main.go
@@ -71,6 +71,7 @@ func main() {
 	// Initialize handlers.
 	ecsInitialize()
 	iasInitialize()
+	routeInitialize()
 
 	// Start the HTTP server.
 	fmt.Printf("Starting HTTP connection (%s)...\nNot using the usual port for HTTP?\nBe sure to use a proxy, otherwise the Wii can't connect!\n", CON.Address)
diff --git a/route.go b/route.go
index ff2ef91..89b2d9c 100644
--- a/route.go
+++ b/route.go
@@ -1,6 +1,8 @@
 package main
 
 import (
+	"crypto/sha256"
+	"database/sql"
 	"fmt"
 	"io/ioutil"
 	"log"
@@ -115,9 +117,15 @@ func (route *Route) Handle() http.Handler {
 			return
 		}
 
-		// IAS-specific actions can have separate values available.
-		// TODO: AUTH
-		// TODO: QUITE A LOT
+		// Check for authentication.
+		if action.NeedsAuthentication {
+			success, err := checkAuthentication(e)
+			// Catch-all in case of invalid formatting or true invalidity.
+			if !success || (err != nil) {
+				http.Error(w, "Unauthorized.", http.StatusUnauthorized)
+				return
+			}
+		}
 
 		// Call this action.
 		action.Callback(e)
@@ -135,6 +143,68 @@ func (route *Route) Handle() http.Handler {
 	})
 }
 
+var routeVerify *sql.Stmt
+
+func routeInitialize() {
+	var err error
+	// What we select does not matter.
+	routeVerify, err = db.Prepare(`SELECT DeviceId FROM userbase WHERE DeviceToken=? AND AccountId=? AND DeviceId=?`)
+	if err != nil {
+		log.Fatalf("route initialize: error preparing statement: %v\n", err)
+	}
+}
+
+// checkAuthentication validates various factors from a given request requiring authentication.
+func checkAuthentication(e *Envelope) (bool, error) {
+	// Get necessary authentication identifiers.
+	deviceToken, err := getKey(e.doc, "DeviceToken")
+	if err != nil {
+		return false, err
+	}
+	accountId, err := getKey(e.doc, "AccountId")
+	if err != nil {
+		return false, err
+	}
+
+	hash := validateTokenFormat(deviceToken)
+	if hash == "" {
+		return false, nil
+	}
+
+	// The Wii sends us a MD5 of its given token. We store a SHA-256 hash of that.
+	hashedDeviceToken := fmt.Sprintf("%x", sha256.Sum256([]byte(hash)))
+
+	// Check using various input given.
+	row := routeVerify.QueryRow(hashedDeviceToken, accountId, e.DeviceId())
+
+	var throwaway string
+	err = row.Scan(&throwaway)
+
+	if err == sql.ErrNoRows {
+		return false, err
+	} else if err != nil {
+		// We shouldn't encounter other errors.
+		log.Printf("error occurred while checking authentication: %v\n", err)
+		return false, err
+	} else {
+		return true, nil
+	}
+}
+
+// validateTokenFormat confirms the prefix and size of tokens,
+// in a format such as WT-5d41402abc4b2a76b9719d911017c592.
+// It returns an empty string on failure.
+func validateTokenFormat(token string) string {
+	success := len(token) == 35 && token[:3] == "WT-"
+
+	if success {
+		// Strips the WT- prefix off.
+		return token[3:35]
+	} else {
+		return ""
+	}
+}
+
 func printError(w http.ResponseWriter, reason string) {
 	http.Error(w, reason, http.StatusInternalServerError)
 	fmt.Println("Failed to handle request: " + reason)