av/protocol/rtsp/client.go

/*
NAME
  client.go

DESCRIPTION
  client.go provides a Client type providing functionality to send RTSP requests
  of methods DESCRIBE, OPTIONS, SETUP and PLAY to an RTSP server.

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"
	"net"
	"net/url"
	"strconv"
)

// Client describes an RTSP Client.
type Client struct {
	cSeq    int
	addr    string
	url     *url.URL
	conn    net.Conn
	session string
}

// NewClient returns a pointer to a new Client. The URL u will be parsed and
// a connection to the RTSP server will be made.
func NewClient(u string) (*Client, error) {
	s := &Client{addr: u}
	var err error
	s.url, err = url.Parse(u)
	if err != nil {
		return nil, err
	}
	s.conn, err = net.Dial("tcp", s.url.Host)
	if err != nil {
		return nil, err
	}
	return s, nil
}

// Describe forms and sends an RTSP request of method DESCRIBE to the RTSP server.
func (s *Client) Describe() (*Response, error) {
	return s.writeRequest(s.url, describe, func(req *Request) { req.Header.Add("Accept", "application/sdp") }, nil)
}

// Options forms and sends an RTSP request of method OPTIONS to the RTSP server.
func (s *Client) Options() (*Response, error) {
	return s.writeRequest(s.url, options, nil, nil)
}

// Setup forms and sends an RTSP request of method SETUP to the RTSP server.
func (s *Client) Setup(track, transport string) (*Response, error) {
	url, err := url.Parse(s.addr + "/" + track)
	if err != nil {
		return nil, err
	}
	return s.writeRequest(
		url,
		setup,
		func(req *Request) {
			req.Header.Add("Transport", transport)
		},
		func(resp *Response) {
			s.session = resp.Header.Get("Session")
		},
	)
}

// Play forms and sends an RTSP request of method PLAY to the RTSP server
func (s *Client) Play() (*Response, error) {
	return s.writeRequest(s.url, play, func(req *Request) { req.Header.Add("Session", s.session) }, nil)
}

// writeRequest creates an RTSP request of the method given and writes it to the
// the Client connection.
//
// headerOp and respOp define any operation that needs
// to occur to the request or response for the given method.
func (c *Client) writeRequest(url *url.URL, method string, headerOp func(*Request), respOp func(*Response)) (*Response, error) {
	req, err := NewRequest(method, c.nextCSeq(), url, nil)
	if err != nil {
		return nil, err
	}

	if headerOp != nil {
		headerOp(req)
	}

	resp, err := c.Do(req)
	if err != nil {
		return nil, err
	}

	if respOp != nil {
		respOp(resp)
	}
	return resp, nil
}

// Do sends the given RTSP request r and reads any responses, return the response
// and any errors.
func (c *Client) Do(r *Request) (*Response, error) {
	err := r.Write(c.conn)
	if err != nil {
		return nil, err
	}

	res, err := ReadResponse(bufio.NewReader(c.conn))
	if err != nil {
		return nil, err
	}

	return res, nil
}

func (s *Client) nextCSeq() string {
	s.cSeq++
	return strconv.Itoa(s.cSeq)
}