Merge pull request #72 from lbryio/feature/remote_ip

Forward real IP in internal-api client calls

extras/lbryinc/client.go

@@ -13,11 +13,29 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
+const (
+	defaultServerAddress = "https://api.lbry.com"
+	timeout              = 5 * time.Second
+	headerForwardedFor   = "X-Forwarded-For"
+
+	userObjectPath = "user"
+	userMeMethod   = "me"
+)
+
 // Client stores data about internal-apis call it is about to make.
 type Client struct {
-	ServerAddress string
 	AuthToken     string
 	Logger        *log.Logger
+	serverAddress string
+	extraHeaders  map[string]string
+}
+
+// ClientOpts allow to provide extra parameters to NewClient:
+// - ServerAddress
+// - RemoteIP — to forward the IP of a frontend client making the request
+type ClientOpts struct {
+	ServerAddress string
+	RemoteIP      string
 }
 
 // APIResponse reflects internal-apis JSON response format.
@@ -30,24 +48,29 @@ type APIResponse struct {
 // ResponseData is a map containing parsed json response.
 type ResponseData map[string]interface{}
 
-const (
-	defaultAPIHost = "https://api.lbry.com"
-	timeout        = 5 * time.Second
-	userObjectPath = "user"
-)
-
 // NewClient returns a client instance for internal-apis. It requires authToken to be provided
 // for authentication.
-func NewClient(authToken string) Client {
-	return Client{
-		ServerAddress: defaultAPIHost,
+func NewClient(authToken string, opts *ClientOpts) Client {
+	c := Client{
+		serverAddress: defaultServerAddress,
+		extraHeaders:  make(map[string]string),
 		AuthToken:     authToken,
 		Logger:        log.StandardLogger(),
 	}
+	if opts != nil {
+		if opts.ServerAddress != "" {
+			c.serverAddress = opts.ServerAddress
+		}
+		if opts.RemoteIP != "" {
+			c.extraHeaders[headerForwardedFor] = opts.RemoteIP
+		}
+	}
+
+	return c
 }
 
 func (c Client) getEndpointURL(object, method string) string {
-	return fmt.Sprintf("%s/%s/%s", c.ServerAddress, object, method)
+	return fmt.Sprintf("%s/%s/%s", c.serverAddress, object, method)
 }
 
 func (c Client) prepareParams(params map[string]interface{}) (string, error) {
@@ -69,9 +92,14 @@ func (c Client) doCall(url string, payload string) ([]byte, error) {
 	if err != nil {
 		return body, err
 	}
+
 	req.Header.Add("Accept", "application/json")
 	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
 
+	for k, v := range c.extraHeaders {
+		req.Header.Set(k, v)
+	}
+
 	client := &http.Client{Timeout: timeout}
 	r, err := client.Do(req)
 	if err != nil {
@@ -107,5 +135,5 @@ func (c Client) Call(object, method string, params map[string]interface{}) (Resp
 
 // UserMe returns user details for the user associated with the current auth_token
 func (c Client) UserMe() (ResponseData, error) {
-	return c.Call(userObjectPath, "me", map[string]interface{}{})
+	return c.Call(userObjectPath, userMeMethod, map[string]interface{}{})
 }

extras/lbryinc/client_test.go

@@ -1,8 +1,8 @@
 package lbryinc
 
 import (
-	"log"
 	"net/http"
+	"net/http/httptest"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -10,19 +10,16 @@ import (
 )
 
 func TestUserMeWrongToken(t *testing.T) {
-	c := NewClient("abc")
+	c := NewClient("abc", nil)
 	r, err := c.UserMe()
 	require.NotNil(t, err)
 	assert.Equal(t, "could not authenticate user", err.Error())
 	assert.Nil(t, r)
 }
 
-const dummyServerURL = "http://127.0.0.1:59999"
-
-func launchDummyServer() {
-	s := &http.Server{
-		Addr: "127.0.0.1:59999",
-		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+func launchDummyServer(lastReq **http.Request) *httptest.Server {
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		*lastReq = &*r
 		w.Header().Set("Content-Type", "application/json; charset=utf-8")
 		w.WriteHeader(http.StatusOK)
 		response := []byte(`{
@@ -50,16 +47,27 @@ func launchDummyServer() {
 			}
 			}`)
 		w.Write(response)
-		}),
-	}
-	log.Fatal(s.ListenAndServe())
+	}))
 }
 
 func TestUserMe(t *testing.T) {
-	go launchDummyServer()
-	c := NewClient("realToken")
-	c.ServerAddress = dummyServerURL
+	var req *http.Request
+	ts := launchDummyServer(&req)
+	defer ts.Close()
+
+	c := NewClient("realToken", &ClientOpts{ServerAddress: ts.URL})
 	r, err := c.UserMe()
 	assert.Nil(t, err)
 	assert.Equal(t, r["primary_email"], "andrey@lbry.com")
 }
+
+func TestRemoteIP(t *testing.T) {
+	var req *http.Request
+	ts := launchDummyServer(&req)
+	defer ts.Close()
+
+	c := NewClient("realToken", &ClientOpts{ServerAddress: ts.URL, RemoteIP: "8.8.8.8"})
+	_, err := c.UserMe()
+	assert.Nil(t, err)
+	assert.Equal(t, []string{"8.8.8.8"}, req.Header["X-Forwarded-For"])
+}