github.com/emcfarlane/larking@v0.0.0-20220605172417-1704b45ee6c3/control/client.go (about)

     1  // Copyright 2021 Edward McFarlane. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package control
     6  
     7  import (
     8  	"context"
     9  	"fmt"
    10  	"io/ioutil"
    11  	"net"
    12  	"net/http"
    13  	"net/url"
    14  	"path"
    15  	"sync"
    16  
    17  	"github.com/emcfarlane/larking/apipb/controlpb"
    18  	"github.com/pkg/browser"
    19  	"golang.org/x/sync/errgroup"
    20  	"google.golang.org/protobuf/encoding/protojson"
    21  )
    22  
    23  // TODO: use OAuth2 libraries directly.
    24  
    25  // Client is the client that connects to a larkingcontrol server for a node.
    26  type Client struct {
    27  	httpc    *http.Client // HTTP client used to talk to larkcontrol
    28  	svrURL   *url.URL     // URL of the larkcontrol server
    29  	cacheDir string       // Cache directory
    30  
    31  	mu     sync.Mutex
    32  	perRPC *PerRPCCredentials
    33  }
    34  
    35  func NewClient(addr, cacheDir string) (*Client, error) {
    36  	svrURL, err := url.Parse(addr)
    37  	if err != nil {
    38  		return nil, err
    39  	}
    40  
    41  	httpc := http.DefaultClient
    42  
    43  	return &Client{
    44  		httpc:    httpc,
    45  		svrURL:   svrURL,
    46  		cacheDir: cacheDir,
    47  	}, nil
    48  
    49  }
    50  
    51  // OpenPerRPCCredentials at the control server address.
    52  func (c *Client) OpenRPCCredentials(ctx context.Context) (*PerRPCCredentials, error) {
    53  	c.mu.Lock()
    54  	defer c.mu.Unlock()
    55  	if c.perRPC != nil {
    56  		return c.perRPC, nil
    57  	}
    58  
    59  	// Login, save creds to file.
    60  	credFile := path.Join(c.cacheDir, "credentials.json")
    61  
    62  	creds, err := c.doLogin(ctx)
    63  	if err != nil {
    64  		return nil, fmt.Errorf("doLogin: %w", err)
    65  	}
    66  
    67  	b, err := protojson.Marshal(creds)
    68  	if err != nil {
    69  		return nil, err
    70  	}
    71  	if err := ioutil.WriteFile(credFile, b, 0644); err != nil {
    72  		return nil, err
    73  	}
    74  
    75  	u := "file://" + credFile
    76  	perRPC, err := OpenRPCCredentials(ctx, u)
    77  	if err != nil {
    78  		return nil, err
    79  	}
    80  	c.perRPC = perRPC
    81  	return perRPC, nil
    82  }
    83  
    84  func (c *Client) doLogin(ctx context.Context) (*controlpb.Credentials, error) {
    85  	// open browser
    86  
    87  	l, err := net.Listen("tcp", "localhost:0")
    88  	if err != nil {
    89  		return nil, err
    90  	}
    91  
    92  	newURL := *c.svrURL
    93  	newURL.Path = path.Join(newURL.Path, "/login")
    94  	values := newURL.Query()
    95  	values.Add("mode", "select")
    96  	values.Add("signInSuccessUrl", "http://"+l.Addr().String())
    97  	newURL.RawQuery = values.Encode()
    98  
    99  	fmt.Println("Opening URL to complete login flow: ", newURL.String())
   100  	if err := browser.OpenURL(newURL.String()); err != nil {
   101  		return nil, err
   102  	}
   103  
   104  	var (
   105  		cred  *controlpb.Credentials
   106  		creds = make(chan *controlpb.Credentials)
   107  	)
   108  
   109  	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   110  		ctx := r.Context()
   111  		if err := func() error {
   112  			//body, err := httputil.DumpRequest(r, true)
   113  			//if err != nil {
   114  			//	return err
   115  			//}
   116  			//fmt.Println("--- body ---")
   117  			//fmt.Println(string(body))
   118  			//fmt.Println()
   119  
   120  			if r.Method != http.MethodGet || r.URL.Path != "/" {
   121  				okURL := *c.svrURL
   122  				okURL.Path = r.URL.Path
   123  				http.Redirect(w, r, okURL.String(), http.StatusFound)
   124  				return nil
   125  			}
   126  
   127  			q := r.URL.Query()
   128  			name := q.Get("name")
   129  			accessToken := q.Get("accessToken")
   130  
   131  			v := controlpb.Credentials{
   132  				Name: name,
   133  				Type: &controlpb.Credentials_Bearer{
   134  					Bearer: &controlpb.Credentials_BearerToken{
   135  						AccessToken: accessToken,
   136  					},
   137  				},
   138  			}
   139  
   140  			select {
   141  			case creds <- &v:
   142  			case <-ctx.Done():
   143  				return ctx.Err()
   144  			}
   145  
   146  			okURL := *c.svrURL
   147  			okURL.Path = path.Join(okURL.Path, "/success")
   148  			http.Redirect(w, r, okURL.String(), http.StatusFound)
   149  			return nil
   150  		}(); err != nil {
   151  			http.Error(w, err.Error(), http.StatusBadRequest)
   152  		}
   153  	})
   154  	server := &http.Server{Handler: handler}
   155  
   156  	g, gctx := errgroup.WithContext(ctx)
   157  	g.Go(func() error {
   158  		if err := server.Serve(l); err != nil {
   159  			if err != http.ErrServerClosed {
   160  				return err
   161  			}
   162  		}
   163  		return nil
   164  	})
   165  	g.Go(func() error {
   166  		select {
   167  		case v := <-creds:
   168  			cred = v
   169  		case <-gctx.Done():
   170  			return gctx.Err()
   171  		}
   172  		return server.Close()
   173  	})
   174  	if err := g.Wait(); err != nil {
   175  		return nil, err
   176  	}
   177  	return cred, nil
   178  }