github.com/emc-advanced-dev/unik@v0.0.0-20190717152701-a58d3e8e33b7/containers/compilers/rump/c/stub/stub.go (about)

     1  package main
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/json"
     6  	"errors"
     7  	"io/ioutil"
     8  	"log"
     9  	"net"
    10  	"net/http"
    11  	"os"
    12  	"strings"
    13  	"time"
    14  )
    15  
    16  // to exit the program with error use fatal method only!!
    17  
    18  var timeout = time.Duration(2 * time.Second)
    19  
    20  func dialTimeout(network, addr string) (net.Conn, error) {
    21  	return net.DialTimeout(network, addr, timeout)
    22  }
    23  
    24  func getEnvAmazon() (map[string]string, error) {
    25  	client := http.Client{
    26  		Transport: &http.Transport{
    27  			Dial: dialTimeout,
    28  		},
    29  	}
    30  	resp, err := client.Get("http://169.254.169.254/latest/user-data")
    31  	if err != nil {
    32  		return nil, err
    33  	}
    34  	data, err := ioutil.ReadAll(resp.Body)
    35  	if err != nil {
    36  		return nil, err
    37  	}
    38  	defer resp.Body.Close()
    39  	var env map[string]string
    40  	if err := json.Unmarshal(data, &env); err != nil {
    41  		return nil, err
    42  	}
    43  	return env, nil
    44  }
    45  
    46  func main() {
    47  	log.SetOutput(os.Stdout)
    48  	log.Printf("unik v0.0 boostrapping beginning...")
    49  	envChan := make(chan map[string]string)
    50  
    51  	closeChan := make(chan struct{})
    52  
    53  	go func() {
    54  		listenerIp, err := getListenerIp(closeChan)
    55  		if err != nil {
    56  			log.Printf("err getting listener ip: %v", err)
    57  			return
    58  		}
    59  		if env, err := registerWithListener(listenerIp); err != nil {
    60  			log.Printf("err registering with listener: %v", err)
    61  			return
    62  		} else {
    63  			envChan <- env
    64  		}
    65  	}()
    66  
    67  	errChan := make(chan error)
    68  	go func() {
    69  		env, err := getEnvAmazon()
    70  		envChan <- env
    71  		errChan <- err
    72  		close(closeChan)
    73  	}()
    74  
    75  envLoop:
    76  	for {
    77  		log.Printf("waiting for UniK bootstrap")
    78  		select {
    79  		case env := <-envChan:
    80  			if env != nil {
    81  				log.Printf("env was set: %v", env)
    82  				setEnv(env)
    83  				break envLoop
    84  			}
    85  		case err := <-errChan:
    86  			log.Printf("error: %v", err)
    87  		}
    88  	}
    89  	log.Printf("continuing to main\n")
    90  }
    91  
    92  func setEnv(env map[string]string) error {
    93  	for key, val := range env {
    94  		os.Setenv(key, val)
    95  	}
    96  	return nil
    97  }
    98  
    99  func registerWithListener(listenerIp string) (map[string]string, error) {
   100  	//get MAC Addr
   101  	ifaces, err := net.Interfaces()
   102  	if err != nil {
   103  		return nil, errors.New("retrieving network interfaces" + err.Error())
   104  	}
   105  	macAddress := ""
   106  	for _, iface := range ifaces {
   107  		log.Printf("found an interface: %v\n", iface)
   108  		if len(iface.HardwareAddr) > 0 {
   109  			macAddress = iface.HardwareAddr.String()
   110  			break
   111  		}
   112  	}
   113  	if macAddress == "" {
   114  		return nil, errors.New("could not find mac address")
   115  	}
   116  
   117  	resp, err := http.Post("http://"+listenerIp+":3000/register?mac_address="+macAddress, "", bytes.NewBuffer([]byte{}))
   118  	if err != nil {
   119  		return nil, err
   120  	}
   121  	data, err := ioutil.ReadAll(resp.Body)
   122  	if err != nil {
   123  		return nil, err
   124  	}
   125  	var env map[string]string
   126  	if err := json.Unmarshal(data, &env); err != nil {
   127  		return nil, err
   128  	}
   129  	return env, nil
   130  }
   131  
   132  func getListenerIp(closeChan <-chan struct{}) (string, error) {
   133  	log.Printf("listening for udp heartbeat...")
   134  	socket, err := net.ListenUDP("udp4", &net.UDPAddr{
   135  		IP:   net.IPv4(0, 0, 0, 0),
   136  		Port: 9967,
   137  	})
   138  	if err != nil {
   139  		return "", err
   140  	}
   141  	for {
   142  		log.Printf("listening...")
   143  		data := make([]byte, 4096)
   144  		_, remoteAddr, err := socket.ReadFromUDP(data)
   145  		if err != nil {
   146  			return "", err
   147  		}
   148  		log.Printf("received an ip from %s with data: %s", remoteAddr.IP.String(), string(data))
   149  		if strings.Contains(string(data), "unik") {
   150  			data = bytes.Trim(data, "\x00")
   151  			return strings.Split(string(data), ":")[1], nil
   152  		}
   153  		select {
   154  		case <-closeChan:
   155  			return "", nil //registered with ec2
   156  		default:
   157  			continue
   158  		}
   159  	}
   160  }