From 90b5346763b2365298751245e1281c8ef74aef76 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Tue, 5 Mar 2024 19:14:02 +0200
Subject: [PATCH] Add timeouts for HTTP requests

---
 libgm/client.go      | 20 ++++++++++++--------
 libgm/http.go        | 10 +++++++---
 libgm/longpoll.go    |  4 +++-
 libgm/pair_google.go |  4 ++--
 4 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/libgm/client.go b/libgm/client.go
index 32804db..ad1b7f3 100644
--- a/libgm/client.go
+++ b/libgm/client.go
@@ -6,6 +6,7 @@ import (
 	"crypto/sha256"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"net/url"
 	"time"
@@ -77,8 +78,9 @@ type Client struct {
 	AuthData *AuthData
 	cfg      *gmproto.Config
 
-	proxy Proxy
-	http  *http.Client
+	httpTransport *http.Transport
+	http          *http.Client
+	lphttp        *http.Client
 }
 
 func NewAuthData() *AuthData {
@@ -92,11 +94,17 @@ func NewClient(authData *AuthData, logger zerolog.Logger) *Client {
 	sessionHandler := &SessionHandler{
 		responseWaiters: make(map[string]chan<- *IncomingRPCMessage),
 	}
+	transport := &http.Transport{
+		DialContext:           (&net.Dialer{Timeout: 10 * time.Second}).DialContext,
+		TLSHandshakeTimeout:   10 * time.Second,
+		ResponseHeaderTimeout: 20 * time.Second,
+	}
 	cli := &Client{
 		AuthData:       authData,
 		Logger:         logger,
 		sessionHandler: sessionHandler,
-		http:           &http.Client{},
+		http:           &http.Client{Transport: transport, Timeout: 2 * time.Minute},
+		lphttp:         &http.Client{Transport: transport, Timeout: 30 * time.Minute},
 
 		pingShortCircuit: make(chan struct{}),
 	}
@@ -127,11 +135,7 @@ func (c *Client) SetProxy(proxy string) error {
 	if err != nil {
 		c.Logger.Fatal().Err(err).Msg("Failed to set proxy")
 	}
-	proxyUrl := http.ProxyURL(proxyParsed)
-	c.http.Transport = &http.Transport{
-		Proxy: proxyUrl,
-	}
-	c.proxy = proxyUrl
+	c.httpTransport.Proxy = http.ProxyURL(proxyParsed)
 	c.Logger.Debug().Any("proxy", proxyParsed.Host).Msg("SetProxy")
 	return nil
 }
diff --git a/libgm/http.go b/libgm/http.go
index d58ac2d..a5af5f0 100644
--- a/libgm/http.go
+++ b/libgm/http.go
@@ -25,10 +25,10 @@ const ContentTypePBLite = "application/json+protobuf"
 
 func (c *Client) makeProtobufHTTPRequest(url string, data proto.Message, contentType string) (*http.Response, error) {
 	ctx := c.Logger.WithContext(context.TODO())
-	return c.makeProtobufHTTPRequestContext(ctx, url, data, contentType)
+	return c.makeProtobufHTTPRequestContext(ctx, url, data, contentType, false)
 }
 
-func (c *Client) makeProtobufHTTPRequestContext(ctx context.Context, url string, data proto.Message, contentType string) (*http.Response, error) {
+func (c *Client) makeProtobufHTTPRequestContext(ctx context.Context, url string, data proto.Message, contentType string, longPoll bool) (*http.Response, error) {
 	var body []byte
 	var err error
 	switch contentType {
@@ -48,7 +48,11 @@ func (c *Client) makeProtobufHTTPRequestContext(ctx context.Context, url string,
 	}
 	util.BuildRelayHeaders(req, contentType, "*/*")
 	c.AddCookieHeaders(req)
-	res, reqErr := c.http.Do(req)
+	client := c.http
+	if longPoll {
+		client = c.lphttp
+	}
+	res, reqErr := client.Do(req)
 	if reqErr != nil {
 		return res, reqErr
 	}
diff --git a/libgm/longpoll.go b/libgm/longpoll.go
index 8685266..409b858 100644
--- a/libgm/longpoll.go
+++ b/libgm/longpoll.go
@@ -3,6 +3,7 @@ package libgm
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
@@ -214,6 +215,7 @@ func (c *Client) doLongPoll(loggedIn bool, onFirstConnect func()) {
 	defer func() {
 		log.Debug().Msg("Long polling stopped")
 	}()
+	ctx := log.WithContext(context.TODO())
 	log.Debug().Str("listen_uuid", listenReqID).Msg("Long polling starting")
 
 	dittoPing := make(chan struct{}, 1)
@@ -253,7 +255,7 @@ func (c *Client) doLongPoll(loggedIn bool, onFirstConnect func()) {
 			url = util.ReceiveMessagesURLGoogle
 			payload.Auth.Network = util.GoogleNetwork
 		}
-		resp, err := c.makeProtobufHTTPRequest(url, payload, ContentTypePBLite)
+		resp, err := c.makeProtobufHTTPRequestContext(ctx, url, payload, ContentTypePBLite, true)
 		if err != nil {
 			if loggedIn {
 				c.triggerEvent(&events.ListenTemporaryError{Error: err})
diff --git a/libgm/pair_google.go b/libgm/pair_google.go
index 260047c..65bfbd2 100644
--- a/libgm/pair_google.go
+++ b/libgm/pair_google.go
@@ -67,7 +67,7 @@ func (c *Client) signInGaiaInitial(ctx context.Context) (*gmproto.SignInGaiaResp
 	payload := c.baseSignInGaiaPayload()
 	payload.UnknownInt3 = 1
 	return typedHTTPResponse[*gmproto.SignInGaiaResponse](
-		c.makeProtobufHTTPRequestContext(ctx, util.SignInGaiaURL, payload, ContentTypePBLite),
+		c.makeProtobufHTTPRequestContext(ctx, util.SignInGaiaURL, payload, ContentTypePBLite, false),
 	)
 }
 
@@ -82,7 +82,7 @@ func (c *Client) signInGaiaGetToken(ctx context.Context) (*gmproto.SignInGaiaRes
 		SomeData: key,
 	}
 	resp, err := typedHTTPResponse[*gmproto.SignInGaiaResponse](
-		c.makeProtobufHTTPRequestContext(ctx, util.SignInGaiaURL, payload, ContentTypePBLite),
+		c.makeProtobufHTTPRequestContext(ctx, util.SignInGaiaURL, payload, ContentTypePBLite, false),
 	)
 	if err != nil {
 		return nil, err