github.com/cloud-foundations/dominator@v0.0.0-20221004181915-6e4fee580046/cmd/hypervisor/control.go (about)

     1  package main
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"os"
     7  	"path/filepath"
     8  
     9  	"github.com/Cloud-Foundations/Dominator/hypervisor/manager"
    10  	"github.com/Cloud-Foundations/Dominator/lib/fsutil"
    11  	"github.com/Cloud-Foundations/Dominator/lib/log"
    12  )
    13  
    14  var shutdownVMsOnNextStop bool
    15  
    16  type flusher interface {
    17  	Flush() error
    18  }
    19  
    20  func acceptControlConnections(m *manager.Manager, listener net.Listener,
    21  	logger log.DebugLogger) {
    22  	for {
    23  		if conn, err := listener.Accept(); err != nil {
    24  			logger.Println(err)
    25  		} else if err := processControlConnection(conn, m, logger); err != nil {
    26  			logger.Println(err)
    27  		}
    28  	}
    29  }
    30  
    31  func configureVMsToStopOnNextStop() {
    32  	sendRequest(connectToControl(), "stop-vms-on-next-stop")
    33  }
    34  
    35  func connectToControl() net.Conn {
    36  	sockAddr := filepath.Join(*stateDir, "control")
    37  	if conn, err := net.Dial("unix", sockAddr); err != nil {
    38  		fmt.Fprintf(os.Stderr, "Error connecting to: %s: %s\n", sockAddr, err)
    39  		os.Exit(1)
    40  		return nil
    41  	} else {
    42  		return conn
    43  	}
    44  }
    45  
    46  func listenForControl(m *manager.Manager, logger log.DebugLogger) error {
    47  	sockAddr := filepath.Join(*stateDir, "control")
    48  	os.Remove(sockAddr)
    49  	if listener, err := net.Listen("unix", sockAddr); err != nil {
    50  		return err
    51  	} else {
    52  		if err := os.Chmod(sockAddr, fsutil.PrivateFilePerms); err != nil {
    53  			return err
    54  		}
    55  		go acceptControlConnections(m, listener, logger)
    56  		return nil
    57  	}
    58  }
    59  
    60  func processControlConnection(conn net.Conn, m *manager.Manager,
    61  	logger log.DebugLogger) error {
    62  	defer conn.Close()
    63  	buffer := make([]byte, 256)
    64  	if nRead, err := conn.Read(buffer); err != nil {
    65  		return fmt.Errorf("error reading request: %s", err)
    66  	} else if nRead < 1 {
    67  		return fmt.Errorf("read short request: %s", err)
    68  	} else {
    69  		request := string(buffer[:nRead])
    70  		if request[nRead-1] != '\n' {
    71  			return fmt.Errorf("request not null-terminated: %s", request)
    72  		}
    73  		request = request[:nRead-1]
    74  		switch request {
    75  		case "stop":
    76  			if _, err := fmt.Fprintln(conn, "ok"); err != nil {
    77  				return err
    78  			}
    79  			os.Remove(m.GetRootCookiePath())
    80  			if shutdownVMsOnNextStop {
    81  				m.ShutdownVMsAndExit()
    82  			} else {
    83  				logger.Println("stopping without shutting down VMs")
    84  				if flusher, ok := logger.(flusher); ok {
    85  					flusher.Flush()
    86  				}
    87  				os.Exit(0)
    88  			}
    89  		case "stop-vms-on-next-stop":
    90  			if _, err := fmt.Fprintln(conn, "ok"); err != nil {
    91  				return err
    92  			}
    93  			shutdownVMsOnNextStop = true
    94  		default:
    95  			if _, err := fmt.Fprintln(conn, "bad request"); err != nil {
    96  				return err
    97  			}
    98  		}
    99  	}
   100  	return nil
   101  }
   102  
   103  func requestStop() {
   104  	sendRequest(connectToControl(), "stop")
   105  }
   106  
   107  func sendRequest(conn net.Conn, request string) {
   108  	if _, err := fmt.Fprintln(conn, request); err != nil {
   109  		fmt.Fprintf(os.Stderr, "Error writing request: %s\n", err)
   110  		os.Exit(1)
   111  	}
   112  	buffer := make([]byte, 256)
   113  	if nRead, err := conn.Read(buffer); err != nil {
   114  		fmt.Fprintf(os.Stderr, "Error reading response: %s\n", err)
   115  		os.Exit(1)
   116  	} else if nRead < 1 {
   117  		fmt.Fprintf(os.Stderr, "Read short response: %s\n", err)
   118  		os.Exit(1)
   119  	} else {
   120  		response := string(buffer[:nRead])
   121  		if response[nRead-1] != '\n' {
   122  			fmt.Fprintf(os.Stderr, "Response not null-terminated: %s\n",
   123  				response)
   124  			os.Exit(1)
   125  		}
   126  		response = response[:nRead-1]
   127  		if response != "ok" {
   128  			fmt.Fprintf(os.Stderr, "Bad response: %s\n", response)
   129  			os.Exit(1)
   130  		} else {
   131  			conn.Read(buffer) // Wait for EOF.
   132  			conn.Close()
   133  			os.Exit(0)
   134  		}
   135  	}
   136  }