github.com/HashDataInc/packer@v1.3.2/packer/rpc/server.go (about)

     1  package rpc
     2  
     3  import (
     4  	"io"
     5  	"log"
     6  	"net/rpc"
     7  
     8  	"github.com/hashicorp/packer/packer"
     9  	"github.com/ugorji/go/codec"
    10  )
    11  
    12  const (
    13  	DefaultArtifactEndpoint      string = "Artifact"
    14  	DefaultBuildEndpoint                = "Build"
    15  	DefaultBuilderEndpoint              = "Builder"
    16  	DefaultCacheEndpoint                = "Cache"
    17  	DefaultCommandEndpoint              = "Command"
    18  	DefaultCommunicatorEndpoint         = "Communicator"
    19  	DefaultHookEndpoint                 = "Hook"
    20  	DefaultPostProcessorEndpoint        = "PostProcessor"
    21  	DefaultProvisionerEndpoint          = "Provisioner"
    22  	DefaultUiEndpoint                   = "Ui"
    23  )
    24  
    25  // Server represents an RPC server for Packer. This must be paired on
    26  // the other side with a Client.
    27  type Server struct {
    28  	mux      *muxBroker
    29  	streamId uint32
    30  	server   *rpc.Server
    31  	closeMux bool
    32  }
    33  
    34  // NewServer returns a new Packer RPC server.
    35  func NewServer(conn io.ReadWriteCloser) *Server {
    36  	mux, _ := newMuxBrokerServer(conn)
    37  	result := newServerWithMux(mux, 0)
    38  	result.closeMux = true
    39  	go mux.Run()
    40  	return result
    41  }
    42  
    43  func newServerWithMux(mux *muxBroker, streamId uint32) *Server {
    44  	return &Server{
    45  		mux:      mux,
    46  		streamId: streamId,
    47  		server:   rpc.NewServer(),
    48  		closeMux: false,
    49  	}
    50  }
    51  
    52  func (s *Server) Close() error {
    53  	if s.closeMux {
    54  		log.Printf("[WARN] Shutting down mux conn in Server")
    55  		return s.mux.Close()
    56  	}
    57  
    58  	return nil
    59  }
    60  
    61  func (s *Server) RegisterArtifact(a packer.Artifact) {
    62  	s.server.RegisterName(DefaultArtifactEndpoint, &ArtifactServer{
    63  		artifact: a,
    64  	})
    65  }
    66  
    67  func (s *Server) RegisterBuild(b packer.Build) {
    68  	s.server.RegisterName(DefaultBuildEndpoint, &BuildServer{
    69  		build: b,
    70  		mux:   s.mux,
    71  	})
    72  }
    73  
    74  func (s *Server) RegisterBuilder(b packer.Builder) {
    75  	s.server.RegisterName(DefaultBuilderEndpoint, &BuilderServer{
    76  		builder: b,
    77  		mux:     s.mux,
    78  	})
    79  }
    80  
    81  func (s *Server) RegisterCache(c packer.Cache) {
    82  	s.server.RegisterName(DefaultCacheEndpoint, &CacheServer{
    83  		cache: c,
    84  	})
    85  }
    86  
    87  func (s *Server) RegisterCommunicator(c packer.Communicator) {
    88  	s.server.RegisterName(DefaultCommunicatorEndpoint, &CommunicatorServer{
    89  		c:   c,
    90  		mux: s.mux,
    91  	})
    92  }
    93  
    94  func (s *Server) RegisterHook(h packer.Hook) {
    95  	s.server.RegisterName(DefaultHookEndpoint, &HookServer{
    96  		hook: h,
    97  		mux:  s.mux,
    98  	})
    99  }
   100  
   101  func (s *Server) RegisterPostProcessor(p packer.PostProcessor) {
   102  	s.server.RegisterName(DefaultPostProcessorEndpoint, &PostProcessorServer{
   103  		mux: s.mux,
   104  		p:   p,
   105  	})
   106  }
   107  
   108  func (s *Server) RegisterProvisioner(p packer.Provisioner) {
   109  	s.server.RegisterName(DefaultProvisionerEndpoint, &ProvisionerServer{
   110  		mux: s.mux,
   111  		p:   p,
   112  	})
   113  }
   114  
   115  func (s *Server) RegisterUi(ui packer.Ui) {
   116  	s.server.RegisterName(DefaultUiEndpoint, &UiServer{
   117  		ui:       ui,
   118  		register: s.server.RegisterName,
   119  	})
   120  }
   121  
   122  // ServeConn serves a single connection over the RPC server. It is up
   123  // to the caller to obtain a proper io.ReadWriteCloser.
   124  func (s *Server) Serve() {
   125  	// Accept a connection on stream ID 0, which is always used for
   126  	// normal client to server connections.
   127  	stream, err := s.mux.Accept(s.streamId)
   128  	if err != nil {
   129  		log.Printf("[ERR] Error retrieving stream for serving: %s", err)
   130  		return
   131  	}
   132  	defer stream.Close()
   133  
   134  	h := &codec.MsgpackHandle{
   135  		RawToString: true,
   136  		WriteExt:    true,
   137  	}
   138  	rpcCodec := codec.GoRpc.ServerCodec(stream, h)
   139  	s.server.ServeCodec(rpcCodec)
   140  }