diff --git a/cli/daemon/daemon.go b/cli/daemon/daemon.go index fed789d..2947037 100644 --- a/cli/daemon/daemon.go +++ b/cli/daemon/daemon.go @@ -88,7 +88,18 @@ func executeConnect(h2c *http2client.Http2Client, cmd *rpc.Command) (string, err func executeGet(h2c *http2client.Http2Client, cmd *rpc.Command) (string, error) { _, includeHeaders := cmd.Options["--include"] - return h2c.Get(cmd.Args[0], includeHeaders) + timeoutString, exists := cmd.Options["--timeout"] + var timeout int + var err error + if exists { + timeout, err = strconv.Atoi(timeoutString) + if err != nil { + return "", fmt.Errorf("%v: invalid timeout", timeoutString) + } + } else { + timeout = 10 + } + return h2c.Get(cmd.Args[0], includeHeaders, timeout) } func executeCommandAndCloseConnection(h2c *http2client.Http2Client, conn net.Conn, sock net.Listener) { diff --git a/http2client/http2client.go b/http2client/http2client.go index 3834202..aa30096 100644 --- a/http2client/http2client.go +++ b/http2client/http2client.go @@ -4,11 +4,13 @@ package http2client import ( "bytes" "crypto/tls" + "errors" "fmt" "github.com/bradfitz/http2" "github.com/bradfitz/http2/hpack" "net" "sort" + "time" ) type Http2Client struct { @@ -134,7 +136,7 @@ func (h2c *Http2Client) nextAvailableStreamId() uint32 { return result } -func (h2c *Http2Client) Get(path string, includeHeaders bool) (string, error) { +func (h2c *Http2Client) Get(path string, includeHeaders bool, timeoutInSeconds int) (string, error) { if h2c.err != nil { return "", h2c.err } @@ -156,23 +158,32 @@ func (h2c *Http2Client) Get(path string, includeHeaders bool) (string, error) { if err != nil { return "", fmt.Errorf("Failed to write HEADERS frame to %v: %v", h2c.host, err.Error()) } - received := <-h2c.streams[streamId].onClosed - // TODO: Check for errors in received headers - headers := "" - if includeHeaders { - sortedHeaderNames := make([]string, len(received.receivedHeaders)) - i := 0 - for name, _ := range received.receivedHeaders { - sortedHeaderNames[i] = name - i++ - } - sort.Strings(sortedHeaderNames) - for _, name := range sortedHeaderNames { - headers += name + ": " + received.receivedHeaders[name] + "\n" + timeout := make(chan bool, 1) + go func() { + time.Sleep(time.Duration(timeoutInSeconds) * time.Second) + timeout <- true + }() + select { + case received := <-h2c.streams[streamId].onClosed: + // TODO: Check for errors in received headers + headers := "" + if includeHeaders { + sortedHeaderNames := make([]string, len(received.receivedHeaders)) + i := 0 + for name, _ := range received.receivedHeaders { + sortedHeaderNames[i] = name + i++ + } + sort.Strings(sortedHeaderNames) + for _, name := range sortedHeaderNames { + headers += name + ": " + received.receivedHeaders[name] + "\n" + } + headers += "\n" } - headers += "\n" + return headers + string(received.receivedData.Bytes()), nil + case <-timeout: + return "", errors.New("Timeout while waiting for response.") } - return headers + string(received.receivedData.Bytes()), nil } func (h2c *Http2Client) isConnected() bool {