reverseproxy - rpc.go

     1  package main
     2  
     3  import (
     4  	"encoding/json"
     5  	"errors"
     6  	"fmt"
     7  	"strconv"
     8  	"sync"
     9  
    10  	"golang.org/x/net/websocket"
    11  	"vimagination.zapto.org/jsonrpc"
    12  )
    13  
    14  const (
    15  	broadcastList = -1 - iota
    16  	broadcastAdd
    17  	broadcastRename
    18  	broadcastRemove
    19  	broadcastAddRedirect
    20  	broadcastAddCommand
    21  	broadcastModifyRedirect
    22  	broadcastModifyCommand
    23  	broadcastRemoveRedirect
    24  	broadcastRemoveCommand
    25  	broadcastStartRedirect
    26  	broadcastStartCommand
    27  	broadcastStopRedirect
    28  	broadcastStopCommand
    29  	broadcastCommandStopped
    30  	broadcastCommandError
    31  )
    32  
    33  type socket struct {
    34  	*jsonrpc.Server
    35  	conn *websocket.Conn
    36  	id   uint64
    37  }
    38  
    39  var (
    40  	connMu sync.Mutex
    41  	conns  = make(map[*socket]struct{})
    42  	nextID uint64
    43  )
    44  
    45  func NewConn(conn *websocket.Conn) {
    46  	s := socket{
    47  		conn: conn,
    48  	}
    49  	s.Server = jsonrpc.New(conn, &s)
    50  
    51  	connMu.Lock()
    52  	nextID++
    53  	s.id = nextID
    54  	conns[&s] = struct{}{}
    55  	connMu.Unlock()
    56  
    57  	s.SendData(buildInitialMessage())
    58  
    59  	s.Handle()
    60  
    61  	connMu.Lock()
    62  	delete(conns, &s)
    63  	connMu.Unlock()
    64  }
    65  
    66  func ShutdownRPC() {
    67  	connMu.Lock()
    68  	for c := range conns {
    69  		c.conn.Close()
    70  	}
    71  	connMu.Unlock()
    72  }
    73  
    74  func (s *socket) HandleRPC(method string, data json.RawMessage) (interface{}, error) {
    75  	switch method {
    76  	case "add":
    77  		return s.add(data)
    78  	case "rename":
    79  		return s.rename(data)
    80  	case "remove":
    81  		return s.remove(data)
    82  	case "addRedirect":
    83  		return s.addRedirect(data)
    84  	case "addCommand":
    85  		return s.addCommand(data)
    86  	case "modifyRedirect":
    87  		return s.modifyRedirect(data)
    88  	case "modifyCommand":
    89  		return s.modifyCommand(data)
    90  	case "removeRedirect":
    91  		return s.removeRedirect(data)
    92  	case "removeCommand":
    93  		return s.removeCommand(data)
    94  	case "startRedirect":
    95  		return s.startRedirect(data)
    96  	case "startCommand":
    97  		return s.startCommand(data)
    98  	case "stopRedirect":
    99  		return s.stopRedirect(data)
   100  	case "stopCommand":
   101  		return s.stopCommand(data)
   102  	case "getCommandPorts":
   103  		return s.getCommandPorts(data)
   104  	}
   105  	return nil, nil
   106  }
   107  
   108  func buildInitialMessage() json.RawMessage {
   109  	config.mu.RLock()
   110  	buf := []byte{'['}
   111  	f := true
   112  	for name, server := range config.Servers {
   113  		if f {
   114  			f = false
   115  		} else {
   116  			buf = append(buf, ',')
   117  		}
   118  		buf = fmt.Appendf(buf, "[%q,[", name)
   119  		first := true
   120  		for id, redirect := range server.Redirects {
   121  			if first {
   122  				first = false
   123  			} else {
   124  				buf = append(buf, ',')
   125  			}
   126  			buf = fmt.Appendf(buf, "[%d,%d,%q,%t,%q,", id, redirect.From, redirect.To, redirect.Start, redirect.err)
   127  			for n, m := range redirect.Match {
   128  				if n > 0 {
   129  					buf = append(buf, ',')
   130  				}
   131  				buf = fmt.Appendf(buf, "[%t,%q]", m.IsSuffix, m.Name)
   132  			}
   133  			buf = append(buf, ']')
   134  		}
   135  		buf = append(buf, ']', ',', '[')
   136  		first = true
   137  		for id, cmd := range server.Commands {
   138  			if first {
   139  				first = false
   140  			} else {
   141  				buf = append(buf, ',')
   142  			}
   143  			buf = fmt.Appendf(buf, "[%d,%q,[", id, cmd.Exe)
   144  			for n, param := range cmd.Params {
   145  				if n > 0 {
   146  					buf = append(buf, ',')
   147  				}
   148  				buf = fmt.Appendf(buf, "%q", param)
   149  			}
   150  			buf = fmt.Appendf(buf, "],%q,{", cmd.WorkDir)
   151  			o := true
   152  			for key, value := range cmd.Env {
   153  				if o {
   154  					o = false
   155  				} else {
   156  					buf = append(buf, ',')
   157  				}
   158  				buf = fmt.Appendf(buf, "%q:%q", key, value)
   159  			}
   160  			buf = fmt.Appendf(buf, "},%d,%q,", cmd.status, cmd.err)
   161  			if cmd.User != nil {
   162  				buf = fmt.Appendf(buf, "{\"uid\":%d,\"gid\":%d}", cmd.User.UID, cmd.User.GID)
   163  			} else {
   164  				buf = append(buf, 'n', 'u', 'l', 'l')
   165  			}
   166  			for _, m := range cmd.Match {
   167  				buf = fmt.Appendf(buf, ",[%t,%q]", m.IsSuffix, m.Name)
   168  			}
   169  			buf = append(buf, ']')
   170  		}
   171  		buf = append(buf, ']', ']')
   172  	}
   173  	buf = append(buf, ']')
   174  	config.mu.RUnlock()
   175  	return buildMessage(-1, json.RawMessage(buf))
   176  }
   177  
   178  const broadcastStart = "{\"id\": -0,\"result\":"
   179  
   180  func buildMessage(id int, data json.RawMessage) json.RawMessage {
   181  	l := len(broadcastStart) + len(data) + 1
   182  	dat := make(json.RawMessage, l)
   183  	copy(dat, broadcastStart)
   184  	copy(dat[len(broadcastStart):], data)
   185  	id = -id
   186  	if id > 9 {
   187  		dat[6] = '-'
   188  		dat[7] = byte('0' + id/10)
   189  	}
   190  	dat[8] = byte('0' + id%10)
   191  	dat[l-1] = '}'
   192  	return dat
   193  }
   194  
   195  func broadcast(id int, data json.RawMessage, except uint64) {
   196  	connMu.Lock()
   197  	var dat json.RawMessage
   198  	for c := range conns {
   199  		if c.id != except {
   200  			if len(dat) == 0 {
   201  				dat = buildMessage(id, data)
   202  			}
   203  			go c.SendData(dat)
   204  		}
   205  	}
   206  	connMu.Unlock()
   207  }
   208  
   209  type nameID struct {
   210  	Server string `json:"server"`
   211  	ID     uint64 `json:"id"`
   212  }
   213  
   214  func (s *socket) getRedirect(n nameID, fn func(*server, *redirect) error) error {
   215  	config.mu.Lock()
   216  	defer config.mu.Unlock()
   217  	serv, ok := config.Servers[n.Server]
   218  	if !ok {
   219  		return ErrNoServer
   220  	}
   221  	r, ok := serv.Redirects[n.ID]
   222  	if !ok {
   223  		return ErrUnknownRedirect
   224  	}
   225  	if r.Start {
   226  		return ErrServerRunning
   227  	}
   228  	if err := fn(serv, r); err != nil {
   229  		return err
   230  	}
   231  	if err := saveConfig(); err != nil {
   232  		return fmt.Errorf("error saving config: %w", err)
   233  	}
   234  	return nil
   235  }
   236  
   237  func (s *socket) getCommand(n nameID, fn func(*server, *command) error) error {
   238  	config.mu.Lock()
   239  	defer config.mu.Unlock()
   240  	serv, ok := config.Servers[n.Server]
   241  	if !ok {
   242  		return ErrNoServer
   243  	}
   244  	c, ok := serv.Commands[n.ID]
   245  	if !ok {
   246  		return ErrUnknownCommand
   247  	}
   248  	if c.status == 1 {
   249  		return ErrServerRunning
   250  	}
   251  	if err := fn(serv, c); err != nil {
   252  		return err
   253  	}
   254  	if err := saveConfig(); err != nil {
   255  		return fmt.Errorf("error saving config: %w", err)
   256  	}
   257  	return nil
   258  }
   259  
   260  func (s *socket) add(data json.RawMessage) (interface{}, error) {
   261  	var name string
   262  	if err := json.Unmarshal(data, &name); err != nil {
   263  		return nil, err
   264  	}
   265  	config.mu.Lock()
   266  	defer config.mu.Unlock()
   267  	if _, ok := config.Servers[name]; ok {
   268  		return nil, ErrNameExists
   269  	}
   270  	config.Servers[name] = &server{
   271  		Redirects: make(map[uint64]*redirect),
   272  		Commands:  make(map[uint64]*command),
   273  		name:      name,
   274  	}
   275  	if err := saveConfig(); err != nil {
   276  		return nil, fmt.Errorf("error saving config: %w", err)
   277  	}
   278  	broadcast(broadcastAdd, data, s.id)
   279  	return nil, nil
   280  }
   281  
   282  func (s *socket) rename(data json.RawMessage) (interface{}, error) {
   283  	var name [2]string
   284  	if err := json.Unmarshal(data, &name); err != nil {
   285  		return nil, err
   286  	}
   287  	config.mu.Lock()
   288  	defer config.mu.Unlock()
   289  	if _, ok := config.Servers[name[1]]; ok {
   290  		return nil, ErrNameExists
   291  	}
   292  	serv, ok := config.Servers[name[0]]
   293  	if !ok {
   294  		return nil, ErrNoServer
   295  	}
   296  	delete(config.Servers, name[0])
   297  	config.Servers[name[1]] = serv
   298  	if err := saveConfig(); err != nil {
   299  		return nil, fmt.Errorf("error saving config: %w", err)
   300  	}
   301  	broadcast(broadcastRename, data, s.id)
   302  	return nil, nil
   303  }
   304  
   305  func (s *socket) remove(data json.RawMessage) (interface{}, error) {
   306  	var name string
   307  	if err := json.Unmarshal(data, &name); err != nil {
   308  		return nil, err
   309  	}
   310  	config.mu.Lock()
   311  	defer config.mu.Unlock()
   312  	serv, ok := config.Servers[name]
   313  	if !ok {
   314  		return nil, ErrNoServer
   315  	}
   316  	for _, r := range serv.Redirects {
   317  		if r.Start {
   318  			return nil, ErrServerRunning
   319  		}
   320  	}
   321  	for _, c := range serv.Commands {
   322  		if c.status != 0 {
   323  			return nil, ErrServerRunning
   324  		}
   325  	}
   326  	delete(config.Servers, name)
   327  	if err := saveConfig(); err != nil {
   328  		return nil, fmt.Errorf("error saving config: %w", err)
   329  	}
   330  	broadcast(broadcastRemove, data, s.id)
   331  	return nil, nil
   332  }
   333  
   334  func (s *socket) addRedirect(data json.RawMessage) (interface{}, error) {
   335  	var ar struct {
   336  		Server string `json:"server"`
   337  		redirectData
   338  	}
   339  	if err := json.Unmarshal(data, &ar); err != nil {
   340  		return nil, err
   341  	}
   342  	config.mu.Lock()
   343  	defer config.mu.Unlock()
   344  	serv, ok := config.Servers[ar.Server]
   345  	if !ok {
   346  		return nil, ErrNoServer
   347  	}
   348  	id := serv.addRedirect(ar.redirectData)
   349  	if err := saveConfig(); err != nil {
   350  		return nil, fmt.Errorf("error saving config: %w", err)
   351  	}
   352  	broadcast(broadcastAddRedirect, append(strconv.AppendUint(append(data[:len(data)-1], ",\"id\":"...), id, 10), '}'), s.id)
   353  	return id, nil
   354  }
   355  
   356  func (s *socket) addCommand(data json.RawMessage) (interface{}, error) {
   357  	var ac struct {
   358  		Server string `json:"server"`
   359  		commandData
   360  	}
   361  	if err := json.Unmarshal(data, &ac); err != nil {
   362  		return nil, err
   363  	}
   364  	config.mu.Lock()
   365  	defer config.mu.Unlock()
   366  	serv, ok := config.Servers[ac.Server]
   367  	if !ok {
   368  		return nil, ErrNoServer
   369  	}
   370  	id := serv.addCommand(ac.commandData)
   371  	if err := saveConfig(); err != nil {
   372  		return nil, fmt.Errorf("error saving config: %w", err)
   373  	}
   374  	broadcast(broadcastAddCommand, append(strconv.AppendUint(append(data[:len(data)-1], ",\"id\":"...), id, 10), '}'), s.id)
   375  	return id, nil
   376  }
   377  
   378  func (s *socket) modifyRedirect(data json.RawMessage) (interface{}, error) {
   379  	var mr struct {
   380  		nameID
   381  		redirectData
   382  	}
   383  	if err := json.Unmarshal(data, &mr); err != nil {
   384  		return nil, err
   385  	}
   386  	return nil, s.getRedirect(mr.nameID, func(_ *server, r *redirect) error {
   387  		r.redirectData = mr.redirectData
   388  		broadcast(broadcastModifyRedirect, data, s.id)
   389  		return nil
   390  	})
   391  }
   392  
   393  func (s *socket) modifyCommand(data json.RawMessage) (interface{}, error) {
   394  	var mc struct {
   395  		nameID
   396  		commandData
   397  	}
   398  	if err := json.Unmarshal(data, &mc); err != nil {
   399  		return nil, err
   400  	}
   401  	return nil, s.getCommand(mc.nameID, func(_ *server, c *command) error {
   402  		c.commandData = mc.commandData
   403  		broadcast(broadcastModifyCommand, data, s.id)
   404  		return nil
   405  	})
   406  }
   407  
   408  func (s *socket) removeRedirect(data json.RawMessage) (interface{}, error) {
   409  	var rr nameID
   410  	if err := json.Unmarshal(data, &rr); err != nil {
   411  		return nil, err
   412  	}
   413  	return nil, s.getRedirect(rr, func(serv *server, _ *redirect) error {
   414  		delete(serv.Redirects, rr.ID)
   415  		broadcast(broadcastRemoveRedirect, data, s.id)
   416  		return nil
   417  	})
   418  }
   419  
   420  func (s *socket) removeCommand(data json.RawMessage) (interface{}, error) {
   421  	var rc nameID
   422  	if err := json.Unmarshal(data, &rc); err != nil {
   423  		return nil, err
   424  	}
   425  	return nil, s.getCommand(rc, func(serv *server, _ *command) error {
   426  		delete(serv.Commands, rc.ID)
   427  		broadcast(broadcastRemoveCommand, data, s.id)
   428  		return nil
   429  	})
   430  }
   431  
   432  func (s *socket) startRedirect(data json.RawMessage) (interface{}, error) {
   433  	var sr nameID
   434  	if err := json.Unmarshal(data, &sr); err != nil {
   435  		return nil, err
   436  	}
   437  	return nil, s.getRedirect(sr, func(_ *server, r *redirect) error {
   438  		r.Run()
   439  		broadcast(broadcastStartRedirect, data, s.id)
   440  		return nil
   441  	})
   442  }
   443  
   444  func (s *socket) startCommand(data json.RawMessage) (interface{}, error) {
   445  	var sc nameID
   446  	if err := json.Unmarshal(data, &sc); err != nil {
   447  		return nil, err
   448  	}
   449  	return nil, s.getCommand(sc, func(_ *server, c *command) error {
   450  		err := c.Run()
   451  		if err != nil {
   452  			broadcast(broadcastCommandError, append(strconv.AppendQuote(append(data[:len(data)-1], ',', '"', 'e', 'r', 'r', '"', ':'), err.Error()), '}'), s.id)
   453  		} else {
   454  			broadcast(broadcastStartCommand, data, s.id)
   455  		}
   456  		return err
   457  	})
   458  }
   459  
   460  func (s *socket) stopRedirect(data json.RawMessage) (interface{}, error) {
   461  	var sr nameID
   462  	if err := json.Unmarshal(data, &sr); err != nil {
   463  		return nil, err
   464  	}
   465  	config.mu.Lock()
   466  	defer config.mu.Unlock()
   467  	serv, ok := config.Servers[sr.Server]
   468  	if !ok {
   469  		return nil, ErrNoServer
   470  	}
   471  	r, ok := serv.Redirects[sr.ID]
   472  	if !ok {
   473  		return nil, ErrUnknownRedirect
   474  	}
   475  	if !r.Start {
   476  		return nil, ErrServerNotRunning
   477  	}
   478  	r.Stop()
   479  	broadcast(broadcastStopRedirect, data, s.id)
   480  	if err := saveConfig(); err != nil {
   481  		return nil, fmt.Errorf("error saving config: %w", err)
   482  	}
   483  	return nil, nil
   484  }
   485  
   486  func (s *socket) stopCommand(data json.RawMessage) (interface{}, error) {
   487  	var sc nameID
   488  	if err := json.Unmarshal(data, &sc); err != nil {
   489  		return nil, err
   490  	}
   491  	config.mu.Lock()
   492  	defer config.mu.Unlock()
   493  	serv, ok := config.Servers[sc.Server]
   494  	if !ok {
   495  		return nil, ErrNoServer
   496  	}
   497  	c, ok := serv.Commands[sc.ID]
   498  	if !ok {
   499  		return nil, ErrUnknownCommand
   500  	}
   501  	if c.status != 1 {
   502  		return nil, ErrServerNotRunning
   503  	}
   504  	c.Stop()
   505  	broadcast(broadcastStopCommand, data, s.id)
   506  	if err := saveConfig(); err != nil {
   507  		return nil, fmt.Errorf("error saving config: %w", err)
   508  	}
   509  	return nil, nil
   510  }
   511  
   512  func (s *socket) getCommandPorts(data json.RawMessage) (interface{}, error) {
   513  	var cp nameID
   514  	if err := json.Unmarshal(data, &cp); err != nil {
   515  		return nil, err
   516  	}
   517  	config.mu.Lock()
   518  	defer config.mu.Unlock()
   519  	serv, ok := config.Servers[cp.Server]
   520  	if !ok {
   521  		return nil, ErrNoServer
   522  	}
   523  	c, ok := serv.Commands[cp.ID]
   524  	if !ok {
   525  		return nil, ErrUnknownCommand
   526  	}
   527  	if c.status != 1 {
   528  		return nil, ErrServerNotRunning
   529  	}
   530  	ports := c.unixCmd.Status().Ports
   531  	return ports, nil
   532  }
   533  
   534  var (
   535  	ErrNameExists       = errors.New("name already exists")
   536  	ErrNoServer         = errors.New("no server by that name exists")
   537  	ErrServerRunning    = errors.New("cannot perform operation while server running")
   538  	ErrServerNotRunning = errors.New("server not running")
   539  	ErrUnknownRedirect  = errors.New("unknown redirect")
   540  	ErrUnknownCommand   = errors.New("unknown command")
   541  )
   542