av/protocol/rtsp/rtsp.go

/*
NAME
  rtsp.go

DESCRIPTION
  rtsp.go provides functionality for forming and sending RTSP requests for
  methods, DESCRIBE, OPTIONS, SETUP and PLAY, as described by
  the RTSP standards, see https://tools.ietf.org/html/rfc7826

AUTHORS
  Saxon A. Nelson-Milton <saxon@ausocean.org>

LICENSE
  This is Copyright (C) 2019 the Australian Ocean Lab (AusOcean).

  It is free software: you can redistribute it and/or modify them
  under the terms of the GNU General Public License as published by the
  Free Software Foundation, either version 3 of the License, or (at your
  option) any later version.

  It is distributed in the hope that it will be useful, but WITHOUT
  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
  for more details.

  You should have received a copy of the GNU General Public License
  in gpl.txt.  If not, see http://www.gnu.org/licenses.
*/
package rtsp

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"
)

// Minimum response size to be considered valid in bytes.
const minResponse = 2

var errSmallResponse = errors.New("response too small")

// Request describes an RTSP request.
type Request struct {
	Method        string
	URL           *url.URL
	Proto         string
	ProtoMajor    int
	ProtoMinor    int
	Header        http.Header
	ContentLength int
	Body          io.ReadCloser
}

// NewRequest returns a pointer to a new Request.
func NewRequest(method, cSeq string, u *url.URL, body io.ReadCloser) (*Request, error) {
	req := &Request{
		Method:     method,
		URL:        u,
		Proto:      "RTSP",
		ProtoMajor: 1,
		ProtoMinor: 0,
		Header:     map[string][]string{"CSeq": []string{cSeq}},
		Body:       body,
	}
	return req, nil
}

// Write writes the request r to the given io.Writer w.
func (r *Request) Write(w io.Writer) error {
	_, err := w.Write([]byte(r.String()))
	return err
}

// String returns a formatted string of the Request.
func (r Request) String() string {
	var b strings.Builder
	b.WriteString(r.Method + " " + r.URL.String() + " " + r.Proto + "/" + strconv.Itoa(r.ProtoMajor) + "." + strconv.Itoa(r.ProtoMinor) + "\r\n")
	for k, v := range r.Header {
		for _, v := range v {
			b.WriteString(k + ": " + v + "\r\n")
		}
	}
	b.WriteString("\r\n")
	if r.Body != nil {
		s, _ := ioutil.ReadAll(r.Body)
		b.WriteString(string(s))
	}
	return b.String()
}

// Response describes an RTSP response.
type Response struct {
	Proto         string
	ProtoMajor    int
	ProtoMinor    int
	StatusCode    int
	Status        string
	ContentLength int
	Header        http.Header
	Body          io.ReadCloser
}

// String returns a formatted string of the Response.
func (r Response) String() string {
	var b strings.Builder
	b.WriteString(r.Proto + "/" + strconv.Itoa(r.ProtoMajor) + "." + strconv.Itoa(r.ProtoMinor) + " " + strconv.Itoa(r.StatusCode) + " " + r.Status + "\n")
	for k, v := range r.Header {
		for _, v := range v {
			b.WriteString(k + ": " + v + "\n")
		}
	}
	return b.String()
}

// ReadResponse will read the response of the RTSP request from the connection,
// and return a pointer to a new Response.
func ReadResponse(r io.Reader) (*Response, error) {
	resp := &Response{Header: make(map[string][]string)}

	b := bufio.NewReader(r)

	// Read the first line.
	s, err := b.ReadString('\n')
	if err != nil {
		return nil, err
	}

	if len(s) < minResponse {
		return nil, errSmallResponse
	}

	if s[:5] != "RTSP/" {
		return nil, errors.New("not a valid RTSP response")
	}
	resp.Proto = "RTSP"

	_, err = fmt.Sscanf(s[5:12], "%d.%d %d", &resp.ProtoMajor, &resp.ProtoMinor, &resp.StatusCode)
	if err != nil {
		return nil, err
	}

	// Read headers.
	for {
		s, err := b.ReadString('\n')
		if err != nil {
			return nil, err
		}

		if strings.TrimRight(s, "\r\n") == "" {
			break
		}

		parts := strings.SplitN(s, ":", 2)
		resp.Header.Add(strings.TrimSpace(parts[0]), strings.TrimSpace(parts[1]))
	}

	// Get the content length from the header.
	resp.ContentLength, err = strconv.Atoi(resp.Header.Get("Content-Length"))
	if err != nil {
		return nil, err
	}

	resp.Body = closer{b, r}
	return resp, nil
}

type closer struct {
	*bufio.Reader
	r io.Reader
}

func (c closer) Close() error {
	if c.Reader == nil {
		return nil
	}
	defer func() {
		c.Reader = nil
		c.r = nil
	}()
	if r, ok := c.r.(io.ReadCloser); ok {
		return r.Close()
	}
	return nil
}