From 5dab3c4e96a97d06cea970f74da87966306e796b Mon Sep 17 00:00:00 2001 From: r4 Date: Mon, 21 Jun 2021 20:37:22 +0200 Subject: [PATCH] Add mp3 support + command line option to limit tracks + general redesign for more modularity --- README.md | 4 +- main.go | 258 +++++++++++++++++++++++++++------------------ model/extractor.go | 17 +++ mp3/mp3.go | 113 ++++++++++++++++++++ util/io.go | 1 + vorbis/vorbis.go | 103 ++++++++---------- 6 files changed, 331 insertions(+), 165 deletions(-) create mode 100644 model/extractor.go create mode 100644 mp3/mp3.go diff --git a/README.md b/README.md index e07b7d4..2f1c114 100644 --- a/README.md +++ b/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. - -MP3 support is planned but not yet implemented. +A program that extracts the individual tracks from an Ogg/Vorbis or mp3 radio stream. Written in go without any non-standard dependencies. ## Building diff --git a/main.go b/main.go index f06ec12..ba5fa6b 100644 --- a/main.go +++ b/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 ( - colRed = "\033[31m" +var client = new(http.Client) + +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...] + `+arg0+` [options...] Options: -dir -- Output directory (default: "."). + -n -- Stop after tracks. Output types: * - ` + colYellow + `! ` + colReset + ` - ` + colRed + `! ` + colReset) + `+colYellow+`! `+colReset+` + `+colRed+`! `+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++ - if i >= len(os.Args) { - printErr("Expected string after flag '%v'", arg) + dir = expectArg(arg) + case "-n": + 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] - case "--help": - usage(os.Args[0], 0) - case "-h": + limitTracks = true + maxTracks = int(n) + case "--help", "-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) - 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 - } - + // Record the actual stream. for { - var raw bytes.Buffer - - // 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 + record(url, dir) + printInfo("Reconnecting due to previous error") } } diff --git a/model/extractor.go b/model/extractor.go new file mode 100644 index 0000000..a900e02 --- /dev/null +++ b/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) +} diff --git a/mp3/mp3.go b/mp3/mp3.go new file mode 100644 index 0000000..a5695b9 --- /dev/null +++ b/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_ + // 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 +} diff --git a/util/io.go b/util/io.go index 031cfce..8d63009 100644 --- a/util/io.go +++ b/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 diff --git a/vorbis/vorbis.go b/vorbis/vorbis.go index 611c94f..8f6c621 100644 --- a/vorbis/vorbis.go +++ b/vorbis/vorbis.go @@ -4,93 +4,78 @@ import ( "bytes" "errors" "io" + "strconv" + "strings" ) var ( - ErrNoHeaderSegment = errors.New("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") + ErrNoHeaderSegment = errors.New("vorbis: no header segment") ) -type Decoder struct { - r io.Reader +type Extractor struct { 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 { - return &Decoder{ - r: r, - } +func NewExtractor() (*Extractor, error) { + return new(Extractor), nil } -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 -// 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 { +func (d *Extractor) TryGetFilename() (filename string, hasFilename bool) { if !d.hasMetadata { - return ErrCallReadRestAfterReadMetadata + return "", false } + d.hasMetadata = false - for { - page, _, err := d.readPage() - if err != nil { - return err - } - - if (page.Header.HeaderType & FHeaderTypeEOS) > 0 { - // End of stream - break + // Use relevant metadata to create a filename. + var base string // Filename without extension. + artist, artistOk := d.metadata.FieldByName("Artist") + title, titleOk := d.metadata.FieldByName("Title") + if artistOk || titleOk { + if !artistOk { + artist = "Unknown" + } else if !titleOk { + 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 }