func (w protobufStreamWriter) Close() error {
_, err := protocol.Messages(&protocol.StreamChunk{
EOF: proto.Bool(true),
}).WriteTo(w.conn)
return err
}
func (s *WardenServer) handleRun(conn net.Conn, request *protocol.RunRequest) (proto.Message, error) {
handle := request.GetHandle()
script := request.GetScript()
privileged := request.GetPrivileged()
env := request.GetEnv()
container, err := s.backend.Lookup(handle)
if err != nil {
return nil, err
}
s.bomberman.Pause(container.Handle())
defer s.bomberman.Unpause(container.Handle())
ProcessSpec := warden.ProcessSpec{
Script: script,
Privileged: privileged,
EnvironmentVariables: convertEnvironmentVariables(env),
}
if request.Rlimits != nil {
ProcessSpec.Limits = resourceLimits(request.Rlimits)
}
processID, stream, err := container.Run(ProcessSpec)
if err != nil {
return nil, err
}
protocol.Messages(&protocol.ProcessPayload{
ProcessId: proto.Uint32(processID),
}).WriteTo(conn)
return s.streamProcessToConnection(processID, stream, conn), nil
}
func (s *WardenServer) streamProcessToConnection(processID uint32, stream <-chan warden.ProcessStream, conn net.Conn) proto.Message {
for payload := range stream {
if payload.ExitStatus != nil {
return &protocol.ProcessPayload{
ProcessId: proto.Uint32(processID),
ExitStatus: proto.Uint32(*payload.ExitStatus),
}
}
var payloadSource protocol.ProcessPayload_Source
switch payload.Source {
case warden.ProcessStreamSourceStdout:
payloadSource = protocol.ProcessPayload_stdout
case warden.ProcessStreamSourceStderr:
payloadSource = protocol.ProcessPayload_stderr
case warden.ProcessStreamSourceStdin:
payloadSource = protocol.ProcessPayload_stdin
}
protocol.Messages(&protocol.ProcessPayload{
ProcessId: proto.Uint32(processID),
Source: &payloadSource,
Data: proto.String(string(payload.Data)),
}).WriteTo(conn)
}
return nil
}
func (s *WardenServer) handleStreamOut(conn net.Conn, request *protocol.StreamOutRequest) (proto.Message, error) {
handle := request.GetHandle()
srcPath := request.GetSrcPath()
container, err := s.backend.Lookup(handle)
if err != nil {
return nil, err
}
s.bomberman.Pause(container.Handle())
defer s.bomberman.Unpause(container.Handle())
_, err = protocol.Messages(&protocol.StreamOutResponse{}).WriteTo(conn)
if err != nil {
return nil, err
}
writer := transport.NewProtobufStreamWriter(conn)
reader, err := container.StreamOut(srcPath)
if err != nil {
return nil, err
}
_, err = io.Copy(writer, reader)
if err != nil {
return nil, err
}
return nil, writer.Close()
}
func (w protobufStreamWriter) Write(buff []byte) (int, error) {
_, err := protocol.Messages(&protocol.StreamChunk{
Content: buff,
}).WriteTo(w.conn)
if err != nil {
return 0, err
}
return len(buff), nil
}
func (s *WardenServer) handleStreamIn(conn net.Conn, reader *bufio.Reader, request *protocol.StreamInRequest) (proto.Message, error) {
handle := request.GetHandle()
dstPath := request.GetDstPath()
container, err := s.backend.Lookup(handle)
if err != nil {
return nil, err
}
s.bomberman.Pause(container.Handle())
defer s.bomberman.Unpause(container.Handle())
streamWriter, err := container.StreamIn(dstPath)
if err != nil {
return nil, err
}
_, err = protocol.Messages(&protocol.StreamInResponse{}).WriteTo(conn)
if err != nil {
return nil, err
}
streamReader := transport.NewProtobufStreamReader(reader)
_, err = io.Copy(streamWriter, streamReader)
if err != nil {
return nil, err
}
err = streamWriter.Close()
if err != nil {
return nil, err
}
return &protocol.StreamInResponse{}, nil
}
"code.google.com/p/gogoprotobuf/proto"
. "github.com/onsi/ginkgo"
. "github.com/onsi/gomega"
protocol "github.com/cloudfoundry-incubator/garden/protocol"
"github.com/cloudfoundry-incubator/garden/transport"
)
var _ = Describe("Reading request messages over the wire", func() {
Context("when a request is received", func() {
It("returns the request and no error", func() {
payload := bufio.NewReader(
protocol.Messages(&protocol.EchoRequest{
Message: proto.String("some-message"),
}),
)
request, err := transport.ReadRequest(payload)
Expect(err).ToNot(HaveOccurred())
Expect(request).To(Equal(
&protocol.EchoRequest{
Message: proto.String("some-message"),
},
))
})
})
Context("when the connection is broken", func() {
It("returns an error", func() {
assertWriteBufferContains := func(messages ...proto.Message) {
reader := bufio.NewReader(bytes.NewBuffer(writeBuffer.Bytes()))
for _, msg := range messages {
req, err := transport.ReadRequest(reader)
Ω(err).ShouldNot(HaveOccurred())
Ω(req).Should(Equal(msg))
}
}
JustBeforeEach(func() {
writeBuffer = bytes.NewBuffer([]byte{})
fakeConn := &FakeConn{
ReadBuffer: protocol.Messages(wardenMessages...),
WriteBuffer: writeBuffer,
}
connection = New(fakeConn)
})
BeforeEach(func() {
wardenMessages = []proto.Message{}
rlimits := &warden.ResourceLimits{
As: proto.Uint64(1),
Core: proto.Uint64(2),
Cpu: proto.Uint64(4),
Data: proto.Uint64(5),
Fsize: proto.Uint64(6),
func (s *WardenServer) serveConnection(conn net.Conn) {
read := bufio.NewReader(conn)
for {
var response proto.Message
var err error
if <-s.stopping {
conn.Close()
break
}
request, err := transport.ReadRequest(read)
if err == io.EOF {
break
}
if err != nil {
log.Println("error reading request:", err)
continue
}
if <-s.stopping {
conn.Close()
break
}
s.openRequests.Incr()
switch req := request.(type) {
case *protocol.PingRequest:
response, err = s.handlePing(req)
case *protocol.EchoRequest:
response, err = s.handleEcho(req)
case *protocol.CapacityRequest:
response, err = s.handleCapacity(req)
case *protocol.CreateRequest:
response, err = s.handleCreate(req)
case *protocol.DestroyRequest:
response, err = s.handleDestroy(req)
case *protocol.ListRequest:
response, err = s.handleList(req)
case *protocol.StopRequest:
response, err = s.handleStop(req)
case *protocol.StreamInRequest:
response, err = s.handleStreamIn(conn, read, req)
case *protocol.StreamOutRequest:
response, err = s.handleStreamOut(conn, req)
case *protocol.RunRequest:
s.openRequests.Decr()
response, err = s.handleRun(conn, req)
s.openRequests.Incr()
case *protocol.AttachRequest:
s.openRequests.Decr()
response, err = s.handleAttach(conn, req)
s.openRequests.Incr()
case *protocol.LimitBandwidthRequest:
response, err = s.handleLimitBandwidth(req)
case *protocol.LimitMemoryRequest:
response, err = s.handleLimitMemory(req)
case *protocol.LimitDiskRequest:
response, err = s.handleLimitDisk(req)
case *protocol.LimitCpuRequest:
response, err = s.handleLimitCpu(req)
case *protocol.NetInRequest:
response, err = s.handleNetIn(req)
case *protocol.NetOutRequest:
response, err = s.handleNetOut(req)
case *protocol.InfoRequest:
response, err = s.handleInfo(req)
default:
err = UnhandledRequestError{request}
}
if err != nil {
response = &protocol.ErrorResponse{
Message: proto.String(err.Error()),
}
}
if response != nil {
protocol.Messages(response).WriteTo(conn)
}
s.openRequests.Decr()
}
}
. "github.com/onsi/ginkgo"
. "github.com/onsi/gomega"
protocol "github.com/cloudfoundry-incubator/garden/protocol"
"github.com/cloudfoundry-incubator/garden/transport"
)
var _ = Describe("Reading response messages over the wire", func() {
Context("when a message of the expected type is received", func() {
It("populates the response object and returns no error", func() {
var echoResponse protocol.EchoResponse
err := transport.ReadMessage(
bufio.NewReader(protocol.Messages(&protocol.EchoRequest{
Message: proto.String("some message"),
})),
&echoResponse,
)
Expect(err).ToNot(HaveOccurred())
Expect(echoResponse.GetMessage()).To(Equal("some message"))
})
})
Context("when the connection is broken", func() {
It("returns an error", func() {
var dummyResponse protocol.PingResponse
payload := protocol.Messages(&protocol.PingRequest{})
JustBeforeEach(func() {
wardenServer = server.New("unix", socketPath, 0, serverBackend)
err := wardenServer.Start()
Expect(err).ToNot(HaveOccurred())
Eventually(ErrorDialing("unix", socketPath)).ShouldNot(HaveOccurred())
serverConnection, err = net.Dial("unix", socketPath)
Expect(err).ToNot(HaveOccurred())
responses = bufio.NewReader(serverConnection)
})
writeMessages := func(message proto.Message) {
num, err := protocol.Messages(message).WriteTo(serverConnection)
Expect(err).ToNot(HaveOccurred())
Expect(num).ToNot(Equal(0))
}
readResponse := func(response proto.Message) {
err := transport.ReadMessage(responses, response)
Expect(err).ToNot(HaveOccurred())
}
It("stops accepting new connections", func() {
go wardenServer.Stop()
Eventually(ErrorDialing("unix", socketPath)).Should(HaveOccurred())
})
It("stops handling requests on existing connections", func() {