Implement authorization middleware

Change WriteJson function to use only for successful response
Implement WriteError function

@@ -10,6 +10,12 @@ import (
 	"github.com/jmoiron/sqlx"
 )
 
+type Prop int
+
+const (
+	PropUserNo Prop = iota
+)
+
 type App struct {
 	Config Config
 	db     *sqlx.DB

 package main
 
 import (
+	"context"
 	"encoding/json"
 	"net/http"
+	"strings"
 	"time"
 
 	"github.com/dgrijalva/jwt-go"
@@ -22,7 +24,7 @@ func (app *App) PostUsers(w http.ResponseWriter, r *http.Request) {
 	body := make(map[string]interface{})
 	err := json.NewDecoder(r.Body).Decode(&body)
 	if err != nil {
-		WriteJson(w, http.StatusBadRequest, map[string]interface{}{"msg": "Failed to parse request json"})
+		WriteError(w, http.StatusBadRequest, "Failed to parse request json")
 		return
 	}
 
@@ -32,48 +34,73 @@ func (app *App) PostUsers(w http.ResponseWriter, r *http.Request) {
 	if err != nil {
 		if merr, ok := err.(*mysql.MySQLError); ok {
 			if merr.Number == 1062 {
-				WriteJson(w, http.StatusConflict, map[string]interface{}{"msg": "Already registered"})
+				WriteError(w, http.StatusConflict, "Already registered")
 				return
 			}
 		}
 
-		WriteJson(w, http.StatusInternalServerError, map[string]interface{}{"msg": "Failed to register"})
+		WriteError(w, http.StatusInternalServerError, "Failed to register")
 		return
 	}
 
 	no, _ := res.LastInsertId()
-	WriteJson(w, http.StatusOK, map[string]interface{}{"user_no": no})
+	WriteJson(w, map[string]interface{}{"user_no": no})
+}
+
+type AuthClaims struct {
+	UserNo uint64 `json:"user_no"`
+	jwt.StandardClaims
 }
 
 func (app *App) PostTokens(w http.ResponseWriter, r *http.Request) {
 	body := make(map[string]interface{})
 	err := json.NewDecoder(r.Body).Decode(&body)
 	if err != nil {
-		WriteJson(w, http.StatusBadRequest, map[string]interface{}{"msg": "Failed to parse request json"})
+		WriteError(w, http.StatusBadRequest, "Failed to parse request json")
 		return
 	}
 
 	hash := sha3.Sum256([]byte(body["password"].(string)))
 	rows, err := app.db.Query("SELECT `no` FROM users WHERE `id`=? AND `password`=?", body["id"], hash[:])
 	if err != nil {
-		WriteJson(w, http.StatusInternalServerError, map[string]interface{}{"msg": "Failed to register"})
+		WriteError(w, http.StatusInternalServerError, "Failed to register")
 		return
 	}
 
 	if !rows.Next() {
-		WriteJson(w, http.StatusUnauthorized, map[string]interface{}{"msg": "Login failed"})
+		WriteError(w, http.StatusUnauthorized, "Login failed")
 		return
 	}
 
 	no := uint64(0)
 	rows.Scan(&no)
 
-	token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{"no": no})
+	token := jwt.NewWithClaims(jwt.SigningMethodHS256, AuthClaims{UserNo: no})
 	auth, err := token.SignedString([]byte(app.Config.TokenSecret))
 	if err != nil {
-		WriteJson(w, http.StatusInternalServerError, map[string]interface{}{"msg": "Login failed"})
+		WriteError(w, http.StatusInternalServerError, "Login failed")
 		return
 	}
 
-	WriteJson(w, http.StatusOK, map[string]interface{}{"token": auth})
+	WriteJson(w, map[string]interface{}{"token": auth})
+}
+
+func (app *App) WithAuth(next func(http.ResponseWriter, *http.Request)) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		auth := r.Header.Get("Authorization")
+		if len(auth) > 6 && strings.Index(auth, "Bearer ") == 0 {
+			token, err := jwt.ParseWithClaims(auth[7:], &AuthClaims{}, func(token *jwt.Token) (interface{}, error) {
+				return []byte(app.Config.TokenSecret), nil
+			})
+
+			if err == nil {
+				claims := token.Claims.(*AuthClaims)
+				ctx := context.WithValue(r.Context(), PropUserNo, claims.UserNo)
+				next(w, r.WithContext(ctx))
+				return
+			}
+		}
+
+		WriteError(w, http.StatusUnauthorized, "Authorization failed")
+	})
 }

@@ -5,8 +5,13 @@ import (
 	"net/http"
 )
 
-func WriteJson(w http.ResponseWriter, status int, data interface{}) {
+func WriteJson(w http.ResponseWriter, data interface{}) {
 	w.Header().Set("Content-Type", "application/json")
-	w.WriteHeader(status)
 	json.NewEncoder(w).Encode(data)
 }
+
+func WriteError(w http.ResponseWriter, status int, message string) {
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(status)
+	json.NewEncoder(w).Encode(map[string]interface{}{"msg": message})
+}