From fd239e4f2193290411a7a84737308e24a85656d8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ferdinand=20M=C3=BCtsch?= <ferdinand@muetsch.io>
Date: Sat, 30 Jan 2021 10:54:54 +0100
Subject: [PATCH] chore: add check to validate wakatime api key before
 accepting it

---
 config/config.go               |  6 +++++
 middlewares/custom/wakatime.go |  7 +-----
 routes/settings.go             | 41 +++++++++++++++++++++++++++++++++-
 3 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/config/config.go b/config/config.go
index 96b8052..f253c74 100644
--- a/config/config.go
+++ b/config/config.go
@@ -33,6 +33,12 @@ const (
 	KeyLatestTotalUsers = "latest_total_users"
 )
 
+const (
+	WakatimeApiUrl                = "https://wakatime.com/api/v1"
+	WakatimeApiHeartbeatsEndpoint = "/users/current/heartbeats.bulk"
+	WakatimeApiUserEndpoint       = "/users/current"
+)
+
 var cfg *Config
 var cFlag = flag.String("config", defaultConfigPath, "config file location")
 
diff --git a/middlewares/custom/wakatime.go b/middlewares/custom/wakatime.go
index eaed9ff..027eb03 100644
--- a/middlewares/custom/wakatime.go
+++ b/middlewares/custom/wakatime.go
@@ -13,11 +13,6 @@ import (
 	"time"
 )
 
-const (
-	WakatimeApiUrl                = "https://wakatime.com/api/v1"
-	WakatimeApiHeartbeatsEndpoint = "/users/current/heartbeats.bulk"
-)
-
 /* Middleware to conditionally relay heartbeats to Wakatime */
 type WakatimeRelayMiddleware struct {
 	httpClient *http.Client
@@ -68,7 +63,7 @@ func (m *WakatimeRelayMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Reque
 
 	go m.send(
 		http.MethodPost,
-		WakatimeApiUrl+WakatimeApiHeartbeatsEndpoint,
+		config.WakatimeApiUrl+config.WakatimeApiHeartbeatsEndpoint,
 		bytes.NewReader(body),
 		headers,
 	)
diff --git a/routes/settings.go b/routes/settings.go
index 4dbc03a..080c2e6 100644
--- a/routes/settings.go
+++ b/routes/settings.go
@@ -1,6 +1,7 @@
 package routes
 
 import (
+	"encoding/base64"
 	"fmt"
 	"github.com/gorilla/mux"
 	"github.com/gorilla/schema"
@@ -12,6 +13,7 @@ import (
 	"log"
 	"net/http"
 	"strconv"
+	"time"
 )
 
 type SettingsHandler struct {
@@ -21,6 +23,7 @@ type SettingsHandler struct {
 	aliasSrvc           services.IAliasService
 	aggregationSrvc     services.IAggregationService
 	languageMappingSrvc services.ILanguageMappingService
+	httpClient          *http.Client
 }
 
 var credentialsDecoder = schema.NewDecoder()
@@ -33,6 +36,7 @@ func NewSettingsHandler(userService services.IUserService, summaryService servic
 		aggregationSrvc:     aggregationService,
 		languageMappingSrvc: languageMappingService,
 		userSrvc:            userService,
+		httpClient:          &http.Client{Timeout: 10 * time.Second},
 	}
 }
 
@@ -255,7 +259,15 @@ func (h *SettingsHandler) PostSetWakatimeApiKey(w http.ResponseWriter, r *http.R
 	}
 
 	user := r.Context().Value(models.UserKey).(*models.User)
-	if _, err := h.userSrvc.SetWakatimeApiKey(user, r.PostFormValue("api_key")); err != nil {
+	apiKey := r.PostFormValue("api_key")
+
+	// Healthcheck, if a new API key is set, i.e. the feature is activated
+	if (user.WakatimeApiKey == "" && apiKey != "") && !h.validateWakatimeKey(apiKey) {
+		templates[conf.SettingsTemplate].Execute(w, h.buildViewModel(r).WithError("failed to connect to WakaTime, API key invalid?"))
+		return
+	}
+
+	if _, err := h.userSrvc.SetWakatimeApiKey(user, apiKey); err != nil {
 		w.WriteHeader(http.StatusInternalServerError)
 		templates[conf.SettingsTemplate].Execute(w, h.buildViewModel(r).WithError("internal server error"))
 		return
@@ -304,6 +316,33 @@ func (h *SettingsHandler) PostRegenerateSummaries(w http.ResponseWriter, r *http
 	templates[conf.SettingsTemplate].Execute(w, h.buildViewModel(r).WithSuccess("summaries are being regenerated – this may take a few seconds"))
 }
 
+func (h *SettingsHandler) validateWakatimeKey(apiKey string) bool {
+	headers := http.Header{
+		"Accept": []string{"application/json"},
+		"Authorization": []string{
+			fmt.Sprintf("Basic %s", base64.StdEncoding.EncodeToString([]byte(apiKey))),
+		},
+	}
+
+	request, err := http.NewRequest(
+		http.MethodGet,
+		conf.WakatimeApiUrl+conf.WakatimeApiUserEndpoint,
+		nil,
+	)
+	if err != nil {
+		return false
+	}
+
+	request.Header = headers
+
+	response, err := h.httpClient.Do(request)
+	if err != nil || response.StatusCode < 200 || response.StatusCode >= 300 {
+		return false
+	}
+
+	return true
+}
+
 func (h *SettingsHandler) buildViewModel(r *http.Request) *view.SettingsViewModel {
 	user := r.Context().Value(models.UserKey).(*models.User)
 	mappings, _ := h.languageMappingSrvc.GetByUser(user.ID)