func (t *task) acceptRequest(conn net.Conn) {
defer t.reqListener.DoneConn(conn)
var respErr error
req := &acomm.Request{}
if err := acomm.UnmarshalConnData(conn, req); err != nil {
respErr = err
}
if err := req.Validate(); err != nil {
respErr = err
}
// Respond to the initial request
resp, err := acomm.NewResponse(req, nil, nil, respErr)
if err != nil {
log.WithFields(log.Fields{
"error": err,
"req": req,
"respErr": respErr,
}).Error("failed to create initial response")
return
}
if err := acomm.SendConnData(conn, resp); err != nil {
return
}
if respErr != nil {
return
}
// Actually perform the task
t.waitgroup.Add(1)
go t.handleRequest(req)
}
// handleRequest runs the task-specific handler and sends the results to the
// request's response hook.
func (t *task) handleRequest(req *acomm.Request) {
defer t.waitgroup.Done()
// Run the task-specific request handler
result, streamAddr, taskErr := t.handler(req)
// Note: The acomm calls log the error already, but we want to have a log
// of the request and response data as well.
resp, err := acomm.NewResponse(req, result, streamAddr, taskErr)
if err != nil {
log.WithFields(log.Fields{
"task": t.name,
"req": req,
"taskResult": result,
"taskErr": taskErr,
"error": err,
}).Error("failed to create response")
return
}
if err := req.Respond(resp); err != nil {
log.WithFields(log.Fields{
"task": t.name,
"req": req,
"taskResult": result,
"taskErr": taskErr,
"error": err,
}).Error("failed to send response")
return
}
}
func (s *ServerSuite) createTaskListener(taskName string, result chan *params) *acomm.UnixListener {
taskListener := acomm.NewUnixListener(filepath.Join(s.configData.SocketDir, taskName, "test.sock"), 0)
if !s.NoError(taskListener.Start(), "failed to start task listener") {
return nil
}
go func() {
for {
conn := taskListener.NextConn()
if conn == nil {
break
}
defer taskListener.DoneConn(conn)
req := &acomm.Request{}
if err := acomm.UnmarshalConnData(conn, req); err != nil {
result <- nil
continue
}
params := ¶ms{}
_ = req.UnmarshalArgs(params)
// Respond to the initial request
resp, _ := acomm.NewResponse(req, nil, nil, nil)
if err := acomm.SendConnData(conn, resp); err != nil {
result <- nil
continue
}
// Response to hook
resp, _ = acomm.NewResponse(req, req.Args, nil, nil)
if err := req.Respond(resp); err != nil {
result <- nil
continue
}
}
}()
time.Sleep(time.Second)
return taskListener
}
func (s *RequestTestSuite) TestHandleResponse() {
sh, eh, handled := generateHandlers()
respErr := errors.New("foobar")
tests := []struct {
description string
sh acomm.ResponseHandler
eh acomm.ResponseHandler
respResult interface{}
respErr error
}{
{"sh handler, err resp", sh, nil, nil, respErr},
{"sh handler, success resp", sh, nil, struct{}{}, nil},
{"eh handler, err resp", nil, eh, nil, respErr},
{"eh handler, success resp", nil, eh, struct{}{}, nil},
{"both handlers, err resp", sh, eh, nil, respErr},
{"both handlers, success resp", sh, eh, struct{}{}, nil},
}
for _, test := range tests {
handled["success"] = 0
handled["error"] = 0
msg := testMsgFunc(test.description)
req, err := acomm.NewRequest("foobar", "unix://foo", struct{}{}, test.sh, test.eh)
if !s.NoError(err, msg("should not fail to build req")) {
continue
}
resp, err := acomm.NewResponse(req, test.respResult, nil, test.respErr)
if !s.NoError(err, msg("should not fail to build resp")) {
continue
}
req.HandleResponse(resp)
if test.respErr != nil {
s.Equal(0, handled["success"], msg("should not have called success handler"))
if test.eh != nil {
s.Equal(1, handled["error"], msg("should have called error handler"))
} else {
s.Equal(0, handled["error"], msg("should not have called error handler"))
}
} else {
if test.sh != nil {
s.Equal(1, handled["success"], msg("should have called success handler"))
} else {
s.Equal(0, handled["success"], msg("should not have called success handler"))
}
s.Equal(0, handled["error"], msg("should not have called error handler"))
}
}
}
func (s *ResponseTestSuite) TestNewResponse() {
result := map[string]string{
"foo": "bar",
}
request, _ := acomm.NewRequest("foobar", "unix://foo", nil, nil, nil)
respErr := errors.New("foobar")
tests := []struct {
description string
request *acomm.Request
result interface{}
err error
expectedErr bool
}{
{"missing request", nil, result, nil, true},
{"missing result and error", request, nil, nil, false},
{"result and error", request, result, respErr, true},
{"result only", request, result, nil, false},
{"error only", request, nil, respErr, false},
}
for _, test := range tests {
msg := testMsgFunc(test.description)
resp, err := acomm.NewResponse(test.request, test.result, nil, test.err)
if test.expectedErr {
s.Error(err, msg("should have failed"))
s.Nil(resp, msg("should not have returned a response"))
} else {
if !s.NoError(err, msg("should have succeeded")) {
s.T().Log(msg(err.Error()))
continue
}
if !s.NotNil(resp, msg("should have returned a response")) {
continue
}
s.Equal(test.request.ID, resp.ID, msg("should have set an ID"))
var result map[string]string
s.NoError(resp.UnmarshalResult(&result))
if test.result == nil {
s.Nil(result, msg("should have nil result"))
} else {
s.Equal(test.result, result, msg("should have set the result"))
}
s.Equal(test.err, resp.Error, msg("should have set the error"))
}
}
}
func (s *Server) acceptInternalRequest(conn net.Conn) {
defer s.internal.DoneConn(conn)
var respErr error
req := &acomm.Request{}
defer func() {
// Respond to the initial request
resp, err := acomm.NewResponse(req, nil, nil, respErr)
if err != nil {
log.WithFields(log.Fields{
"error": err,
"req": req,
"respErr": respErr,
}).Error("failed to create initial response")
return
}
if err := acomm.SendConnData(conn, resp); err != nil {
log.WithFields(log.Fields{
"error": err,
"req": req,
"respErr": respErr,
}).Error("failed to create initial response")
return
}
}()
if err := acomm.UnmarshalConnData(conn, req); err != nil {
respErr = err
return
}
if err := req.Validate(); err != nil {
respErr = err
return
}
if err := s.handleRequest(req); err != nil {
respErr = err
return
}
return
}
func (s *TrackerTestSuite) TestProxyUnix() {
unixReq, err := s.Tracker.ProxyUnix(s.Request, 0)
s.Error(err, "should fail to proxy when tracker is not listening")
s.Nil(unixReq, "should not return a request")
if !s.NoError(s.Tracker.Start(), "listner should start") {
return
}
unixReq, err = s.Tracker.ProxyUnix(s.Request, 0)
s.NoError(err, "should not fail proxying when tracker is listening")
s.NotNil(unixReq, "should return a request")
s.Equal(s.Request.ID, unixReq.ID, "new request should share ID with original")
s.Equal("unix", unixReq.ResponseHook.Scheme, "new request should have a unix response hook")
s.Equal(1, s.Tracker.NumRequests(), "should have tracked the new request")
resp, err := acomm.NewResponse(unixReq, struct{}{}, nil, nil)
if !s.NoError(err, "new response should not error") {
return
}
if !s.NoError(acomm.Send(unixReq.ResponseHook, resp), "response send should not error") {
return
}
lastResp := s.NextResp()
if !s.NotNil(lastResp, "response should have been proxied to original http response hook") {
return
}
s.Equal(resp.ID, lastResp.ID, "response should have been proxied to original http response hook")
s.Equal(0, s.Tracker.NumRequests(), "should have removed the request from tracking")
// Should not proxy a request already using unix response hook
origUnixReq, err := acomm.NewRequest("foobar", "unix://foo", struct{}{}, nil, nil)
if !s.NoError(err, "new request shoudl not error") {
return
}
unixReq, err = s.Tracker.ProxyUnix(origUnixReq, 0)
s.NoError(err, "should not error with unix response hook")
s.Equal(origUnixReq, unixReq, "should not proxy unix response hook")
s.Equal(0, s.Tracker.NumRequests(), "should not response an unproxied request")
}
// externalHandler is the http handler for external requests.
func (s *Server) externalHandler(w http.ResponseWriter, r *http.Request) {
var respErr error
req := &acomm.Request{}
// Send the immediate response
defer func() {
resp, err := acomm.NewResponse(req, nil, nil, respErr)
respJSON, err := json.Marshal(resp)
if err != nil {
log.WithFields(log.Fields{
"error": err,
"req": req,
"response": resp,
}).Error("failed to marshal initial response")
}
if _, err := w.Write(respJSON); err != nil {
log.WithFields(log.Fields{
"error": err,
"req": req,
"response": resp,
}).Error("failed to send initial response")
}
}()
// Parse the request
body, err := ioutil.ReadAll(r.Body)
if err != nil {
respErr = err
return
}
if err := json.Unmarshal(body, req); err != nil {
respErr = err
return
}
respErr = s.handleRequest(req)
}