Add mp3 support + command line option to limit tracks + general redesign for more modularity

README.md

@@ -1,8 +1,6 @@
 # radio-stream-recorder
 
-A program that extracts the individual tracks from an Ogg/Vorbis http radio stream (without any loss in quality). Written in go without any non-standard dependencies.
+A program that extracts the individual tracks from an Ogg/Vorbis or mp3 radio stream. Written in go without any non-standard dependencies.
-
-MP3 support is planned but not yet implemented.
 
 ## Building
 

main.go

@@ -3,48 +3,56 @@ package main
 import (
 	"bytes"
 	"fmt"
-	"io"
 	"net/http"
 	"os"
 	"path"
 	"strconv"
-	"strings"
-	"time"
 
+	"rsr/model"
+	"rsr/mp3"
 	"rsr/util"
 	"rsr/vorbis"
 )
 
-var (
+var client = new(http.Client)
-	colRed = "\033[31m"
+
+const (
+	colRed    = "\033[31m"
 	colYellow = "\033[33m"
-	colReset = "\033[m"
+	colReset  = "\033[m"
+)
+
+var (
+	nTracksRecorded int // Number of recorded tracks.
+	limitTracks bool
+	maxTracks int
 )
 
 func usage(arg0 string, exitStatus int) {
 	fmt.Fprintln(os.Stderr, `Usage:
-  ` + arg0 + ` [options...] <STREAM_URL>
+  `+arg0+` [options...] <STREAM_URL>
 
 Options:
   -dir <DIRECTORY>  --  Output directory (default: ".").
+  -n <NUM>          --  Stop after <NUM> tracks.
 
 Output types:
   * <INFO>
-  ` + colYellow + `! <WARNING>` + colReset + `
+  `+colYellow+`! <WARNING>`+colReset+`
-  ` + colRed + `! <ERROR>` + colReset)
+  `+colRed+`! <ERROR>`+colReset)
 	os.Exit(exitStatus)
 }
 
 func printInfo(f string, v ...interface{}) {
-	fmt.Printf("* " + f + "\n", v...)
+	fmt.Printf("* "+f+"\n", v...)
 }
 
 func printWarn(f string, v ...interface{}) {
-	fmt.Fprintf(os.Stderr, colYellow + "! " + f + colReset + "\n", v...)
+	fmt.Fprintf(os.Stderr, colYellow+"! "+f+colReset+"\n", v...)
 }
 
 func printNonFatalErr(f string, v ...interface{}) {
-	fmt.Fprintf(os.Stderr, colRed + "! " + f + colReset + "\n", v...)
+	fmt.Fprintf(os.Stderr, colRed+"! "+f+colReset+"\n", v...)
 }
 
 func printErr(f string, v ...interface{}) {
@@ -52,6 +60,115 @@ func printErr(f string, v ...interface{}) {
 	os.Exit(1)
 }
 
+func record(url, dir string) {
+	req, err := http.NewRequest("GET", url, nil)
+	if err != nil {
+		printErr("HTTP request error: %v", err)
+	}
+	req.Header.Add("Icy-MetaData", "1") // Request metadata for icecast mp3 streams.
+	resp, err := client.Do(req)
+	if err != nil {
+		printErr("HTTP error: %v", err)
+	}
+	defer resp.Body.Close()
+
+	var extractor model.Extractor
+
+	// Set up extractor depending on content type.
+	contentType := resp.Header.Get("content-type")
+	supported := "Ogg/Vorbis ('application/ogg'), mp3 ('audio/mpeg')"
+	err = nil
+	switch contentType {
+	case "application/ogg":
+		extractor, err = vorbis.NewExtractor()
+	case "audio/mpeg":
+		extractor, err = mp3.NewExtractor(resp.Header)
+	default:
+		printErr("Content type '%v' not supported, supported formats: %v", contentType, supported)
+	}
+	if err != nil {
+		printErr("%v", err)
+	}
+
+	printInfo("Stream type: '%v'", contentType)
+
+	// Make reader blocking.
+	r := util.NewWaitReader(resp.Body)
+
+	// The first track is always discarded, as streams usually don't start at
+	// the exact end of a track, meaning it is almost certainly going to be
+	// incomplete.
+	discard := true
+
+	var rawFile bytes.Buffer
+	var filename string
+	var hasFilename bool
+
+	for {
+		var block bytes.Buffer
+
+		wasFirst, err := extractor.ReadBlock(r, &block)
+		if err != nil {
+			printNonFatalErr("Error reading block: %v", err)
+			// Reconnect, because this error is usually caused by a
+			// file corruption or a network error.
+			return
+		}
+
+		if wasFirst &&
+			// We only care about the beginning of a new file when it marks an
+			// old file's end, which is not the case in the beginning of the
+			// first file.
+			rawFile.Len() > 0 {
+			if !discard {
+				// Save previous track.
+				if !hasFilename {
+					printNonFatalErr("Error: Could not get a track filename")
+					continue
+				}
+				filePath := path.Join(dir, filename)
+				err := os.WriteFile(filePath, rawFile.Bytes(), 0666)
+				if err != nil {
+					printNonFatalErr("Error writing file: %v", err)
+					continue
+				}
+				printInfo("Saved track as: %v", filePath)
+
+				// Stop after the defined number of tracks (if the option was
+				// given).
+				nTracksRecorded++
+				if limitTracks && nTracksRecorded >= maxTracks {
+					printInfo("Successfully recorded %v tracks, exiting", nTracksRecorded)
+					os.Exit(0)
+				}
+			} else {
+				// See declaration of `discard`.
+				discard = false
+			}
+
+			// Reset everything.
+			rawFile.Reset()
+			hasFilename = false
+		}
+
+		// Try to find out the current track's filename.
+		if !hasFilename {
+			if f, ok := extractor.TryGetFilename(); ok {
+				if discard {
+					printInfo("Discarding track: %v", f)
+				} else {
+					printInfo("Recording track: %v", f)
+				}
+				filename = f
+				hasFilename = true
+			}
+		}
+
+		// Append block to the current file byte buffer.
+		rawFile.Write(block.Bytes())
+	}
+}
+
 func main() {
 	var url string
 	dir := "."
@@ -62,27 +179,39 @@ func main() {
 
 	// Parse command line arguments.
 	for i := 1; i < len(os.Args); i++ {
+		// Returns the argument after the given option. Errors if there is no
+		// argument.
+		expectArg := func(currArg string) string {
+			i++
+			if i >= len(os.Args) {
+				printErr("Expected argument after option '%v'", currArg)
+			}
+			return os.Args[i]
+		}
+
 		arg := os.Args[i]
 		if len(arg) >= 1 && arg[0] == '-' {
-			switch(arg) {
+			switch arg {
 			case "-dir":
-				i++
+				dir = expectArg(arg)
-				if i >= len(os.Args) {
+			case "-n":
-					printErr("Expected string after flag '%v'", arg)
+				nStr := expectArg(arg)
+				n, err := strconv.ParseInt(nStr, 10, 32)
+				if err != nil || n <= 0 {
+					printErr("'%v' is not an integer larger than zero", nStr)
 				}
-				dir = os.Args[i]
+				limitTracks = true
-			case "--help":
+				maxTracks = int(n)
-				usage(os.Args[0], 0)
+			case "--help", "-h":
-			case "-h":
 				usage(os.Args[0], 0)
 			default:
-				printErr("Unknown flag: '%v'", arg)
+				printErr("Unknown option: '%v'", arg)
 			}
 		} else {
 			if url == "" {
 				url = arg
 			} else {
-				printErr("Expected flag, but got '%v'", arg)
+				printErr("Expected option, but got '%v'", arg)
 			}
 		}
 	}
@@ -94,88 +223,11 @@ func main() {
 
 	printInfo("URL: %v", url)
 	printInfo("Output directory: %v", dir)
+	printInfo("Stopping after %v tracks", maxTracks)
 
-	resp, err := http.Get(url)
+	// Record the actual stream.
-	if err != nil {
-		printErr("HTTP error: %v", err)
-	}
-	defer resp.Body.Close()
-
-	contentType := resp.Header.Get("content-type")
-	if contentType != "application/ogg" {
-		printErr("Expected content type 'application/ogg', but got: '%v'", contentType)
-	}
-
-	waitReader := util.NewWaitReader(resp.Body)
-
-	// The first track is always discarded, as streams usually don't start at
-	// the exact end of a track, meaning it is almost certainly going to be
-	// incomplete.
-	discard := true
-
-	printErrWhileRecording := func(f string, v ...interface{}) {
-		printNonFatalErr(f, v...)
-		printWarn("Unable to download track, skipping.")
-		discard = true
-	}
-
 	for {
-		var raw bytes.Buffer
+		record(url, dir)
-
+		printInfo("Reconnecting due to previous error")
-		// Write all the bytes of the stream we'll read into a buffer to be able
-		// save it to a file later.
-		r := io.TeeReader(waitReader, &raw)
-
-		d := vorbis.NewDecoder(r)
-
-		// Read until metadata of the track. Keep in mind that the read bytes
-		// are also copied to the buffer `raw` because of the tee reader.
-		md, checksum, err := d.ReadMetadata()
-		if err != nil {
-			printErrWhileRecording("Error reading metadata: %v", err)
-			printInfo("Retrying in 1s")
-			time.Sleep(1 * time.Second)
-			continue
-		}
-
-		// Create filename based on the extracted metadata
-		var base string // File name without path or extension.
-		artist, artistOk := md.FieldByName("Artist")
-		title, titleOk := md.FieldByName("Title")
-		if artistOk || titleOk {
-			base = artist + " -- " + title
-		} else {
-			base = "Unknown_" + strconv.FormatInt(int64(checksum), 10)
-		}
-		base = strings.ReplaceAll(base, "/", "_") // Replace invalid characters
-
-		if discard {
-			printInfo("Going to discard incomplete track: %v", base)
-		} else {
-			printInfo("Recording track: %v", base)
-		}
-
-		filename := path.Join(dir, base+".ogg")
-
-		// Determine the (extent of) the rest of the track by reading it, saving
-		// the exact contents of the single track to our buffer `raw` using the
-		// tee reader we set up previously.
-		err = d.ReadRest()
-		if err != nil {
-			printErrWhileRecording("Error reading stream: %v", err)
-			continue
-		}
-
-		// See declaration of `discard`.
-		if !discard {
-			err := os.WriteFile(filename, raw.Bytes(), 0666)
-			if err != nil {
-				printErrWhileRecording("Error writing file: %v", err)
-				continue
-			}
-			printInfo("Saved track as: %v", filename)
-		}
-
-		discard = false
 	}
 }

model/extractor.go

@@ -0,0 +1,17 @@
+package model
+
+import (
+	"io"
+)
+
+type Extractor interface {
+	// Reads a single "block" from a radio stream. A block can be any chunk of
+	// data, depending on the file format, for example in Ogg/Vorbis it would
+	// be equivalent to a chunk. Writes the part containing the actual music
+	// data into `w`.
+	// `isFirst` is true, if the block read was the first block of a file.
+	ReadBlock(r io.Reader, w io.Writer) (isFirst bool, err error)
+	// Potentially returns a filename using format-specific metadata. Usually
+	// available after the first few blocks of a file were read.
+	TryGetFilename() (filename string, hasFilename bool)
+}

mp3/mp3.go

@@ -0,0 +1,113 @@
+package mp3
+
+import (
+	"bytes"
+	"encoding/binary"
+	"errors"
+	"hash/crc32"
+	"html"
+	"io"
+	"net/http"
+	"strconv"
+	"strings"
+)
+
+var (
+	ErrNoMetaint         = errors.New("mp3: key 'icy-metaint' not found in HTTP header")
+	ErrCorruptedMetadata = errors.New("mp3: corrupted metadata")
+	ErrNoStreamTitle     = errors.New("mp3: no 'StreamTitle' tag in metadata")
+)
+
+type Extractor struct {
+	metaint        int64 // Distance between two metadata chunks
+	hasStreamTitle bool
+	streamTitle    string // Metadata tag determining the filename
+}
+
+func NewExtractor(respHdr http.Header) (*Extractor, error) {
+	mi := respHdr.Get("icy-metaint")
+	if mi == "" {
+		return nil, ErrNoMetaint
+	}
+	miNum, _ := strconv.ParseInt(mi, 10, 64)
+	return &Extractor{
+		metaint: miNum,
+	}, nil
+}
+
+func (d *Extractor) ReadBlock(r io.Reader, w io.Writer) (isFirst bool, err error) {
+	var musicData bytes.Buffer
+
+	// We want to write everything to the output, as well as musicData for
+	// calculating the checksum.
+	multi := io.MultiWriter(w, &musicData)
+
+	// Read until the metadata chunk. The part that is read here is also what
+	// contains the actual mp3 music data.
+	io.CopyN(multi, r, d.metaint)
+
+	// Read number of metadata blocks (blocks within this function are not what
+	// is meant with `ReadBlock()`).
+	var numBlocks uint8
+	err = binary.Read(r, binary.LittleEndian, &numBlocks)
+
+	// Whether this block is the beginning of a new track.
+	var isBOF bool
+
+	// Read metadata blocks.
+	if numBlocks > 0 {
+		// Metadata is only actually stored in the first metadata chunk
+		// of a given file. Therefore, every metadata chunk with more than 1
+		// block always marks the beginning of a file.
+		isBOF = true
+
+		// Each block is 16 bytes in size. Any excess bytes in the last block
+		// are set to '\0', which is great because the `string()` conversion
+		// function ignores null bytes. The whole string is escaped via HTML.
+		// Metadata format: k0='v0';k1='v1';
+		raw := make([]byte, numBlocks*16)
+		if _, err := r.Read(raw); err != nil {
+			return false, err
+		}
+		rawString := html.UnescapeString(string(raw))
+		for _, data := range strings.Split(rawString, ";") {
+			s := strings.Split(data, "=")
+			if len(s) == 2 {
+				if s[0] == "StreamTitle" {
+					d.hasStreamTitle = true
+					// Strip stream title's first and last character (single
+					// quotes).
+					t := s[1]
+					if len(t) < 2 {
+						return false, ErrCorruptedMetadata
+					}
+					t = t[1 : len(t)-1]
+					if t == "Unknown" {
+						// If there is no stream title, use format:
+						// Unknown_<crc32 checksum>
+						// Where the checksum is only that of the first block.
+						sumStr := strconv.FormatInt(int64(crc32.ChecksumIEEE(musicData.Bytes())), 10)
+						d.streamTitle = "Unknown_" + sumStr
+					} else {
+						d.streamTitle = t
+					}
+				}
+			} else if len(s) != 1 {
+				return false, ErrCorruptedMetadata
+			}
+		}
+		if !d.hasStreamTitle {
+			return false, ErrNoStreamTitle
+		}
+	}
+
+	return isBOF, nil
+}
+
+func (d *Extractor) TryGetFilename() (filename string, hasFilename bool) {
+	if !d.hasStreamTitle {
+		return "", false
+	}
+	base := strings.ReplaceAll(d.streamTitle, "/", "_") // Replace invalid characters.
+	return base + ".mp3", true
+}

util/io.go

@@ -23,6 +23,7 @@ func (r WaitReader) Read(p []byte) (int, error) {
 		// Reattempt to read the unread bytes until we can fill `p` completely
 		// (or an error occurs).
 		nNew, err := r.r.Read(p[n:len(p)])
+		// Add new number of read bytes.
 		n += nNew
 		if err != nil {
 			return n, err

vorbis/vorbis.go

@@ -4,93 +4,78 @@ import (
 	"bytes"
 	"errors"
 	"io"
+	"strconv"
+	"strings"
 )
 
 var (
-	ErrNoHeaderSegment               = errors.New("no header segment")
+	ErrNoHeaderSegment = errors.New("vorbis: no header segment")
-	ErrNoMetadata                    = errors.New("no metadata found")
-	ErrCallReadRestAfterReadMetadata = errors.New("please call vorbis.Decoder.ReadRest() after having called vorbis.Decoder.ReadMetadata()")
-	ErrReadMetadataCalledTwice       = errors.New("cannot call vorbis.Decoder.ReadMetadata() twice on the same file")
 )
 
-type Decoder struct {
+type Extractor struct {
-	r           io.Reader
 	hasMetadata bool
+	metadata    *VorbisComment // Used for filename.
+	checksum    uint32         // Used for an alternate filename when there's no metadata.
 }
 
-func NewDecoder(r io.Reader) *Decoder {
+func NewExtractor() (*Extractor, error) {
-	return &Decoder{
+	return new(Extractor), nil
-		r: r,
-	}
 }
 
-func (d *Decoder) readPage() (page OggPage, hdr VorbisHeader, err error) {
+func (d *Extractor) ReadBlock(reader io.Reader, w io.Writer) (isFirst bool, err error) {
+	// Everything we read here is part of the music data so we can just use a
+	// tee reader.
+	r := io.TeeReader(reader, w)
+
 	// Decode page.
-	page, err = OggDecode(d.r)
+	page, err := OggDecode(r)
 	if err != nil {
-		return page, hdr, err
+		return false, err
 	}
 
 	// We need to be able to access `page.Segments[0]`.
 	if len(page.Segments) == 0 {
-		return page, hdr, ErrNoHeaderSegment
+		return false, ErrNoHeaderSegment
 	}
 
 	// Decode Vorbis header, stored in `page.Segments[0]`.
-	hdr, err = VorbisHeaderDecode(bytes.NewBuffer(page.Segments[0]))
+	hdr, err := VorbisHeaderDecode(bytes.NewBuffer(page.Segments[0]))
 	if err != nil {
-		return page, hdr, err
+		return false, err
 	}
 
-	return page, hdr, nil
+	// Extract potential metadata.
+	if hdr.PackType == PackTypeComment {
+		d.hasMetadata = true
+		d.metadata = hdr.Comment
+		d.checksum = page.Header.Checksum
+	}
+
+	// Return true for isFirst if this block is the beginning of a new file.
+	return (page.Header.HeaderType & FHeaderTypeBOS) > 0, nil
 }
 
-// Reads the Ogg/Vorbis file until it finds its metadata. Leaves the reader
+func (d *Extractor) TryGetFilename() (filename string, hasFilename bool) {
-// right after the end of the metadata. `crc32Sum` gives the crc32 checksum
-// of the page containing the metadata. It is equivalent to the page checksum
-// used in the Ogg container. Since the page contains more than just metadata,
-// the checksum can usually be used as a unique identifier.
-func (d *Decoder) ReadMetadata() (metadata *VorbisComment, crc32Sum uint32, err error) {
-	if d.hasMetadata {
-		return nil, 0, ErrReadMetadataCalledTwice
-	}
-
-	for {
-		page, hdr, err := d.readPage()
-		if err != nil {
-			return nil, 0, err
-		}
-
-		if (page.Header.HeaderType & FHeaderTypeEOS) > 0 {
-			// End of stream
-			return nil, 0, ErrNoMetadata
-		}
-
-		if hdr.PackType == PackTypeComment {
-			d.hasMetadata = true
-			return hdr.Comment, page.Header.Checksum, nil
-		}
-	}
-}
-
-// Must to be called after `ReadMetadata()`. Reads the rest of the Ogg/Vorbis
-// file, leaving the reader right after the end of the Ogg/Vorbis file.
-func (d *Decoder) ReadRest() error {
 	if !d.hasMetadata {
-		return ErrCallReadRestAfterReadMetadata
+		return "", false
 	}
+	d.hasMetadata = false
 
-	for {
+	// Use relevant metadata to create a filename.
-		page, _, err := d.readPage()
+	var base string // Filename without extension.
-		if err != nil {
+	artist, artistOk := d.metadata.FieldByName("Artist")
-			return err
+	title, titleOk := d.metadata.FieldByName("Title")
-		}
+	if artistOk || titleOk {
-
+		if !artistOk {
-		if (page.Header.HeaderType & FHeaderTypeEOS) > 0 {
+			artist = "Unknown"
-			// End of stream
+		} else if !titleOk {
-			break
+			title = "Unknown"
 		}
+		base = artist + " -- " + title
+	} else {
+		base = "Unknown_" + strconv.FormatInt(int64(d.checksum), 10)
 	}
+	base = strings.ReplaceAll(base, "/", "_") // Replace invalid characters.
 
-	return nil
+	return base + ".ogg", true
 }