github.com/jlowellwofford/u-root@v1.0.0/pkg/sos/client.go (about)

     1  // Copyright 2018 the u-root Authors. 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 sos
     6  
     7  import (
     8  	"bytes"
     9  	"context"
    10  	"encoding/json"
    11  	"fmt"
    12  	"net"
    13  	"net/http"
    14  	"os"
    15  	"os/signal"
    16  	"strconv"
    17  	"strings"
    18  	"syscall"
    19  
    20  	"github.com/gorilla/mux"
    21  )
    22  
    23  // RegistersNecessaryPatterns registers all the necessary patterns needed
    24  // to make a service becomes a SoS client.
    25  func RegistersNecessaryPatterns(router *mux.Router) {
    26  	router.HandleFunc("/ping", func(w http.ResponseWriter, r *http.Request) {
    27  		w.Write([]byte("pong"))
    28  	}).Methods("GET")
    29  }
    30  
    31  // RegisterServiceWithSos tries to register a service with SoS.
    32  // If an non-nil error is returned, the service needs to exit immediately.
    33  func RegisterServiceWithSos(service string, port uint) error {
    34  	return registerServiceWithSos(service, port, "http://localhost:"+PortNum)
    35  }
    36  
    37  func registerServiceWithSos(service string, port uint, sosServerURL string) error {
    38  	m := RegisterReqJson{service, port}
    39  	return makeRequestToServer("POST", sosServerURL+"/register", m)
    40  }
    41  
    42  // UnregisterServiceWithSos makes a request to SoS Server to unregister the service.
    43  // This function should be called before a service exit.
    44  func UnregisterServiceWithSos(service string) error {
    45  	return unregisterServiceWithSos(service, "http://localhost:"+PortNum)
    46  }
    47  
    48  func unregisterServiceWithSos(service string, sosServerURL string) error {
    49  	m := UnRegisterReqJson{service}
    50  	return makeRequestToServer("POST", sosServerURL+"/unregister", m)
    51  }
    52  
    53  func makeRequestToServer(reqType, url string, reqJson interface{}) error {
    54  	b, err := json.Marshal(reqJson)
    55  	if err != nil {
    56  		return err
    57  	}
    58  
    59  	req, err := http.NewRequest(reqType, url, bytes.NewBuffer(b))
    60  	if err != nil {
    61  		return err
    62  	}
    63  
    64  	res, err := http.DefaultClient.Do(req)
    65  	if err != nil {
    66  		return err
    67  	}
    68  
    69  	if res.StatusCode < 200 || res.StatusCode >= 300 {
    70  		decoder := json.NewDecoder(res.Body)
    71  		defer res.Body.Close()
    72  		var retMsg struct{ Error string }
    73  		if err := decoder.Decode(&retMsg); err != nil {
    74  			return err
    75  		}
    76  		if retMsg.Error != "" {
    77  			return fmt.Errorf(retMsg.Error)
    78  		}
    79  	}
    80  
    81  	return nil
    82  }
    83  
    84  // GetListener starts listener on a random port in localhost
    85  // and returns the listener and the port that the listener is on.
    86  // Remember to close the listener with:
    87  //
    88  // defer listener.close()
    89  //
    90  // or if it's used in a server, remember to shutdown the server.
    91  func GetListener() (net.Listener, uint, error) {
    92  	listener, err := net.Listen("tcp", "localhost:0")
    93  	if err != nil {
    94  		return nil, 0, err
    95  	}
    96  
    97  	addrSplit := strings.Split(listener.Addr().String(), ":")
    98  	if len(addrSplit) != 2 {
    99  		listener.Close()
   100  		return nil, 0, fmt.Errorf("Address format not recognized: %v", listener.Addr().String())
   101  	}
   102  
   103  	port, err := strconv.ParseUint(addrSplit[1], 10, 32)
   104  	if err != nil {
   105  		listener.Close()
   106  		return nil, 0, err
   107  	}
   108  	return listener, uint(port), nil
   109  }
   110  
   111  // StartServiceServer establishes registers all necessary patterns to the router passed in,
   112  // registers the service with SoS using the port passed in, and starts serving the service on the
   113  // listener passed in. If any of the above step fails, this function will return an error.
   114  // This function wraps around RegistersNecessaryPatterns, RegisterServiceWithSos, and UnregisterServiceWithSos.
   115  // If no extenral settings are required, instead of calling each of the above separately, one can call
   116  // the GetListener function and pass the result into this function to start and serve their HTTP server right away.
   117  func StartServiceServer(router *mux.Router, serviceName string, listener net.Listener, port uint) error {
   118  	RegistersNecessaryPatterns(router)
   119  	if err := RegisterServiceWithSos(serviceName, port); err != nil {
   120  		return err
   121  	}
   122  	defer UnregisterServiceWithSos(serviceName)
   123  
   124  	shutdownChan := make(chan bool, 2)
   125  	server := http.Server{Handler: router}
   126  	defer func() {
   127  		shutdownChan <- true
   128  	}() // Use to collect any other failure besides signals
   129  
   130  	// Signals Collector
   131  	sigs := make(chan os.Signal, 1)
   132  	signal.Notify(sigs, syscall.SIGHUP, syscall.SIGINT, syscall.SIGQUIT)
   133  	go func() {
   134  		sig := <-sigs
   135  		fmt.Printf("Received: %v\n", sig)
   136  		shutdownChan <- true
   137  	}()
   138  
   139  	// Server Shutdown code
   140  	go func() {
   141  		<-shutdownChan
   142  		fmt.Println("Shutting down...")
   143  		server.Shutdown(context.Background())
   144  	}()
   145  
   146  	return server.Serve(listener)
   147  }