feat: clientmanager has checkconnection

pkg/pmapi/check_connection_test.go

@@ -0,0 +1,91 @@
+// Copyright (c) 2020 Proton Technologies AG
+//
+// This file is part of ProtonMail Bridge.
+//
+// ProtonMail Bridge is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// ProtonMail Bridge is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with ProtonMail Bridge.  If not, see <https://www.gnu.org/licenses/>.
+
+package pmapi
+
+import (
+	"net/http"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/ProtonMail/proton-bridge/pkg/dialer"
+	"github.com/stretchr/testify/require"
+)
+
+const testServerPort = "18000"
+const testRequestTimeout = 10 * time.Second
+
+func TestMain(m *testing.M) {
+	go startServer()
+	time.Sleep(100 * time.Millisecond) // We need to wait till server is fully running.
+	code := m.Run()
+	os.Exit(code)
+}
+
+func startServer() {
+	http.HandleFunc("/ok", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte("OK"))
+	})
+	http.HandleFunc("/timeout", func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(10 * time.Second)
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte("OK"))
+	})
+	http.HandleFunc("/serverError", func(w http.ResponseWriter, r *http.Request) {
+		http.Error(w, "error", http.StatusInternalServerError)
+	})
+	panic(http.ListenAndServe(":"+testServerPort, nil))
+}
+
+func TestCheckConnection(t *testing.T) {
+	checkCheckConnection(t, "ok", "")
+}
+
+func TestCheckConnectionTimeout(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping test in short mode.")
+	}
+
+	checkCheckConnection(t, "timeout", "Client.Timeout exceeded while awaiting headers")
+}
+
+func TestCheckConnectionServerError(t *testing.T) {
+	checkCheckConnection(t, "serverError", "HTTP status code 500")
+}
+
+func checkCheckConnection(t *testing.T, path string, expectedErrMessage string) {
+	client := dialer.DialTimeoutClient()
+	client.Timeout = testRequestTimeout
+
+	ch := make(chan error)
+
+	go checkConnection(client, "http://localhost:"+testServerPort+"/"+path, ch)
+
+	timeout := time.After(testRequestTimeout + time.Second)
+	select {
+	case err := <-ch:
+		if expectedErrMessage == "" {
+			require.NoError(t, err)
+		} else {
+			require.Error(t, err, expectedErrMessage)
+		}
+	case <-timeout:
+		t.Error("checkConnection timeout failed")
+	}
+}

pkg/pmapi/clientmanager.go

@@ -255,6 +255,68 @@ func (cm *ClientManager) GetClientAuthChannel() chan ClientAuth {
 	return cm.clientAuths
 }
 
+// Errors for possible connection issues
+var (
+	ErrNoInternetConnection = errors.New("no internet connection")
+	ErrCanNotReachAPI       = errors.New("can not reach PM API")
+)
+
+// CheckConnection returns an error if there is no internet connection.
+// This should be moved to the ConnectionManager when it is implemented.
+func (cm *ClientManager) CheckConnection() error {
+	client := getHTTPClient(cm.config, cm.roundTripper)
+
+	// Do not cumulate timeouts, use goroutines.
+	retStatus := make(chan error)
+	retAPI := make(chan error)
+
+	// Check protonstatus.com without SSL for performance reasons. vpn_status endpoint is fast and
+	// returns only OK; this endpoint is not known by the public. We check the connection only.
+	go checkConnection(client, "http://protonstatus.com/vpn_status", retStatus)
+
+	// Check of API reachability also uses a fast endpoint.
+	go checkConnection(client, cm.GetRootURL()+"/tests/ping", retAPI)
+
+	errStatus := <-retStatus
+	errAPI := <-retAPI
+
+	switch {
+	case errStatus == nil && errAPI == nil:
+		return nil
+
+	case errStatus == nil && errAPI != nil:
+		cm.log.Error("ProtonStatus is reachable but API is not")
+		return ErrCanNotReachAPI
+
+	case errStatus != nil && errAPI == nil:
+		cm.log.Warn("API is reachable but protonstatus is not")
+		return nil
+
+	case errStatus != nil && errAPI != nil:
+		cm.log.Error("Both ProtonStatus and API are unreachable")
+		return ErrNoInternetConnection
+	}
+
+	return nil
+}
+
+func checkConnection(client *http.Client, url string, errorChannel chan error) {
+	resp, err := client.Get(url)
+	if err != nil {
+		errorChannel <- err
+		return
+	}
+
+	_ = resp.Body.Close()
+
+	if resp.StatusCode != 200 {
+		errorChannel <- fmt.Errorf("HTTP status code %d", resp.StatusCode)
+		return
+	}
+
+	errorChannel <- nil
+}
+
 // forwardClientAuths handles all incoming auths from clients before forwarding them on the bridge auth channel.
 func (cm *ClientManager) forwardClientAuths() {
 	for auth := range cm.clientAuths {
@@ -266,7 +328,7 @@ func (cm *ClientManager) forwardClientAuths() {
 }
 
 // SetTokenIfUnset sets the token for the given userID if it wasn't already set.
-// The token does not expire.
+// The set token does not expire.
 func (cm *ClientManager) SetTokenIfUnset(userID, token string) {
 	cm.tokensLocker.Lock()
 	defer cm.tokensLocker.Unlock()
@@ -352,6 +414,7 @@ func (cm *ClientManager) handleClientAuth(ca ClientAuth) {
 	cm.setToken(ca.UserID, ca.Auth.GenToken(), time.Duration(ca.Auth.ExpiresIn)*time.Second)
 }
 
+// watchTokenExpirations refreshes any tokens which are about to expire.
 func (cm *ClientManager) watchTokenExpirations() {
 	for userID := range cm.expiredTokens {
 		log := cm.log.WithField("userID", userID)