github.com/daniellockard/packer@v0.7.6-0.20141210173435-5a9390934716/packer/rpc/server.go (about)

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