diff --git a/config/config.go b/config/config.go index a24047f..30c092a 100644 --- a/config/config.go +++ b/config/config.go @@ -215,7 +215,7 @@ func readVersion() string { log.Fatal(err) } - return string(bytes) + return strings.TrimSpace(string(bytes)) } func readLanguageColors() map[string]string { diff --git a/main.go b/main.go index 9a83615..7604e6b 100644 --- a/main.go +++ b/main.go @@ -13,6 +13,7 @@ import ( "github.com/gorilla/mux" "github.com/muety/wakapi/middlewares" + customMiddleware "github.com/muety/wakapi/middlewares/custom" "github.com/muety/wakapi/routes" shieldsV1Routes "github.com/muety/wakapi/routes/compat/shields/v1" wtV1Routes "github.com/muety/wakapi/routes/compat/wakatime/v1" @@ -138,6 +139,7 @@ func main() { userService, []string{"/api/health", "/api/compat/shields/v1"}, ).Handler + wakatimeRelayMiddleware := customMiddleware.NewWakatimeRelayMiddleware().Handler // Router configs router.Use(loggingMiddleware, recoveryMiddleware) @@ -169,10 +171,13 @@ func main() { settingsRouter.Path("/regenerate").Methods(http.MethodPost).HandlerFunc(settingsHandler.PostRegenerateSummaries) // API Routes - apiRouter.Path("/heartbeat").Methods(http.MethodPost).HandlerFunc(heartbeatHandler.ApiPost) apiRouter.Path("/summary").Methods(http.MethodGet).HandlerFunc(summaryHandler.ApiGet) apiRouter.Path("/health").Methods(http.MethodGet).HandlerFunc(healthHandler.ApiGet) + heartbeatsApiRouter := apiRouter.Path("/heartbeat").Methods(http.MethodPost).Subrouter() + heartbeatsApiRouter.Use(wakatimeRelayMiddleware) + heartbeatsApiRouter.Path("").HandlerFunc(heartbeatHandler.ApiPost) + // Wakatime compat V1 API Routes wakatimeV1Router.Path("/users/{user}/all_time_since_today").Methods(http.MethodGet).HandlerFunc(wakatimeV1AllHandler.ApiGet) wakatimeV1Router.Path("/users/{user}/summaries").Methods(http.MethodGet).HandlerFunc(wakatimeV1SummariesHandler.ApiGet) diff --git a/middlewares/custom/wakatime.go b/middlewares/custom/wakatime.go new file mode 100644 index 0000000..ec659bf --- /dev/null +++ b/middlewares/custom/wakatime.go @@ -0,0 +1,97 @@ +package relay + +import ( + "bytes" + "encoding/base64" + "fmt" + "github.com/muety/wakapi/config" + "github.com/muety/wakapi/models" + "io" + "io/ioutil" + "log" + "net/http" + "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 +} + +func NewWakatimeRelayMiddleware() *WakatimeRelayMiddleware { + return &WakatimeRelayMiddleware{ + httpClient: &http.Client{ + Timeout: 10 * time.Second, + }, + } +} + +func (m *WakatimeRelayMiddleware) Handler(h http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + m.ServeHTTP(w, r, h.ServeHTTP) + }) +} + +func (m *WakatimeRelayMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) { + defer next(w, r) + + if r.Method != http.MethodPost { + return + } + + user := r.Context().Value(models.UserKey).(*models.User) + if user == nil || user.WakatimeApiKey == "" { + return + } + + body, _ := ioutil.ReadAll(r.Body) + r.Body.Close() + r.Body = ioutil.NopCloser(bytes.NewBuffer(body)) + + headers := http.Header{ + "X-Machine-Name": r.Header.Values("X-Machine-Name"), + "Content-Type": r.Header.Values("Content-Type"), + "Accept": r.Header.Values("Accept"), + "User-Agent": r.Header.Values("User-Agent"), + "X-Origin": []string{ + fmt.Sprintf("wakapi v%s", config.Get().Version), + }, + "Authorization": []string{ + fmt.Sprintf("Basic %s", base64.StdEncoding.EncodeToString([]byte(user.WakatimeApiKey))), + }, + } + + go m.send( + http.MethodPost, + WakatimeApiUrl+WakatimeApiHeartbeatsEndpoint, + bytes.NewReader(body), + headers, + ) +} + +func (m *WakatimeRelayMiddleware) send(method, url string, body io.Reader, headers http.Header) { + request, err := http.NewRequest(method, url, body) + if err != nil { + log.Printf("error constructing relayed request – %v\n", err) + } + + for k, v := range headers { + for _, h := range v { + request.Header.Set(k, h) + } + } + + response, err := m.httpClient.Do(request) + if err != nil { + log.Printf("error executing relayed request – %v\n", err) + } + + if response.StatusCode < 200 || response.StatusCode >= 300 { + log.Printf("failed to relay request, got status %d\n", response.StatusCode) + } +} diff --git a/models/user.go b/models/user.go index 9fcd8e4..50aff5f 100644 --- a/models/user.go +++ b/models/user.go @@ -7,6 +7,7 @@ type User struct { CreatedAt CustomTime `gorm:"type:timestamp; default:CURRENT_TIMESTAMP"` LastLoggedInAt CustomTime `gorm:"type:timestamp; default:CURRENT_TIMESTAMP"` BadgesEnabled bool `json:"-" gorm:"default:false; type:bool"` + WakatimeApiKey string `json:"-"` } type Login struct {