Implemented rewinding for OGG files

main/src/main/java/com/sedmelluq/discord/lavaplayer/container/ogg/OggAudioTrack.java

@@ -1,16 +1,13 @@
 package com.sedmelluq.discord.lavaplayer.container.ogg;
 
-import com.sedmelluq.discord.lavaplayer.tools.FriendlyException;
 import com.sedmelluq.discord.lavaplayer.tools.io.SeekableInputStream;
 import com.sedmelluq.discord.lavaplayer.track.AudioTrackInfo;
 import com.sedmelluq.discord.lavaplayer.track.BaseAudioTrack;
-import com.sedmelluq.discord.lavaplayer.track.playback.AudioProcessingContext;
 import com.sedmelluq.discord.lavaplayer.track.playback.LocalAudioTrackExecutor;
-import java.io.IOException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import static com.sedmelluq.discord.lavaplayer.tools.FriendlyException.Severity.SUSPICIOUS;
+import java.io.IOException;
 
 /**
  * Audio track which handles an OGG stream.
@@ -31,34 +28,19 @@ public OggAudioTrack(AudioTrackInfo trackInfo, SeekableInputStream inputStream)
   }
 
   @Override
-  public void process(final LocalAudioTrackExecutor localExecutor) {
+  public void process(final LocalAudioTrackExecutor localExecutor) throws IOException {
     OggPacketInputStream packetInputStream = new OggPacketInputStream(inputStream, false);
-
-    log.debug("Starting to play an OGG stream track {}", getIdentifier());
-
-    localExecutor.executeProcessingLoop(() -> {
-      try {
-        processTrackLoop(packetInputStream, localExecutor.getProcessingContext());
-      } catch (IOException e) {
-        throw new FriendlyException("Stream broke when playing OGG track.", SUSPICIOUS, e);
-      }
-    }, null, true);
-  }
-
-  private void processTrackLoop(OggPacketInputStream packetInputStream, AudioProcessingContext context) throws IOException, InterruptedException {
     OggTrackBlueprint blueprint = OggTrackLoader.loadTrackBlueprint(packetInputStream);
 
     if (blueprint == null) {
       throw new IOException("Stream terminated before the first packet.");
     }
 
-    while (blueprint != null) {
-      try (OggTrackHandler handler = blueprint.loadTrackHandler(packetInputStream)) {
-        handler.initialise(context, 0, 0);
-        handler.provideFrames();
-      }
+    OggTrackHandler handler = blueprint.loadTrackHandler(packetInputStream);
 
-      blueprint = OggTrackLoader.loadTrackBlueprint(packetInputStream);
-    }
+    log.debug("Starting to play an OGG track {}", getIdentifier());
+
+    handler.initialise(localExecutor.getProcessingContext(), 0, 0);
+    localExecutor.executeProcessingLoop(handler::provideFrames, handler::seekToTimecode, true);
   }
-}
+}

main/src/main/java/com/sedmelluq/discord/lavaplayer/container/ogg/OggContainerProbe.java

@@ -38,7 +38,7 @@ public MediaContainerDetectionResult probe(AudioReference reference, SeekableInp
       return null;
     }
 
-    log.debug("Track {} is an OGG stream.", reference.identifier);
+    log.debug("Track {} is an OGG file.", reference.identifier);
 
     AudioTrackInfoBuilder infoBuilder = AudioTrackInfoBuilder.create(reference, stream).setIsStream(true);
 
@@ -61,6 +61,7 @@ private void collectStreamInformation(SeekableInputStream stream, AudioTrackInfo
     OggMetadata metadata = OggTrackLoader.loadMetadata(packetInputStream);
 
     if (metadata != null) {
+      if (metadata.getLength() != null) infoBuilder.setIsStream(false);
       infoBuilder.apply(metadata);
     }
   }

main/src/main/java/com/sedmelluq/discord/lavaplayer/container/ogg/OggMetadata.java

@@ -15,12 +15,14 @@ public class OggMetadata implements AudioTrackInfoProvider {
   private static final String ARTIST_FIELD = "ARTIST";
 
   private final Map<String, String> tags;
+  private final long length;
 
   /**
    * @param tags Map of OGG metadata with OGG-specific keys.
    */
   public OggMetadata(Map<String, String> tags, Long length) {
     this.tags = tags;
+    this.length = length;
   }
 
   @Override
@@ -35,7 +37,7 @@ public String getAuthor() {
 
   @Override
   public Long getLength() {
-    return null;
+    return length;
   }
 
   @Override

main/src/main/java/com/sedmelluq/discord/lavaplayer/container/ogg/OggPacketInputStream.java

@@ -2,11 +2,13 @@
 
 import com.sedmelluq.discord.lavaplayer.tools.io.SeekableInputStream;
 import com.sedmelluq.discord.lavaplayer.tools.io.StreamTools;
+
 import java.io.DataInput;
 import java.io.DataInputStream;
 import java.io.EOFException;
 import java.io.IOException;
 import java.io.InputStream;
+import java.util.List;
 
 import static com.sedmelluq.discord.lavaplayer.container.MediaContainerDetection.checkNextBytes;
 
@@ -26,6 +28,7 @@ public class OggPacketInputStream extends InputStream {
   private final DataInput dataInput;
   private final int[] segmentSizes;
 
+  private List<OggSeekPoint> seekPoints;
   private OggPageHeader pageHeader;
   private int bytesLeftInPacket;
   private boolean packetContinues;
@@ -44,6 +47,10 @@ public OggPacketInputStream(SeekableInputStream inputStream, boolean closeDelega
     this.state = State.TRACK_BOUNDARY;
   }
 
+  public void setSeekPoints(List<OggSeekPoint> seekPoints) {
+    this.seekPoints = seekPoints;
+  }
+
   /**
    * Load the next track from the stream. This is only valid when the stream is in a track boundary state.
    * @return True if next track is present in the stream, false if the stream has terminated.
@@ -69,6 +76,8 @@ public boolean startNewTrack() {
   public boolean startNewPacket() throws IOException {
     if (state == State.TRACK_BOUNDARY) {
       return false;
+    } else if (state == State.TRACK_SEEKING) {
+      loadNextNonEmptyPage();
     } else if (state != State.PACKET_BOUNDARY) {
       throw new IllegalStateException("Cannot start a new packet while the previous one has not been consumed.");
     }
@@ -274,6 +283,62 @@ public void close() throws IOException {
     }
   }
 
+  /**
+   * Seeks the stream to the specified timecode.
+   * @param timecode Timecode in milliseconds to seek to.
+   * @return The actual timecode in milliseconds to which the stream was seeked.
+   * @throws IOException On read error.
+   */
+  public long seek(long timecode) throws IOException {
+    if (seekPoints == null) {
+      throw new IllegalStateException("Seek points have not been set.");
+    }
+
+    // Binary search for the seek point with the largest timecode that is smaller than or equal to the target timecode
+    int low = 0;
+    int mid = 0;
+    int high = seekPoints.size() - 1;
+    while (low <= high) {
+      mid = (low + high) / 2;
+      if (seekPoints.get(mid).getTimecode() <= timecode) {
+        low = mid + 1;
+      } else {
+        high = mid - 1;
+      }
+    }
+
+    if (mid > 0) {
+      mid--;
+    } else {
+      mid++;
+    }
+
+    OggSeekPoint seekPoint = seekPoints.get(mid);
+    inputStream.seek(seekPoint.getPosition());
+    state = State.TRACK_SEEKING;
+
+    return seekPoint.getTimecode();
+  }
+
+  public List<OggSeekPoint> createSeekTable(int sampleRate) throws IOException {
+    if (!inputStream.canSeekHard()) {
+      return null;
+    }
+
+    long savedPosition = inputStream.getPosition();
+
+    long absoluteOffset = pageHeader.byteStreamPosition;
+    inputStream.seek(absoluteOffset);
+
+    byte[] data = new byte[(int) inputStream.getContentLength()];
+    int dataLength = StreamTools.readUntilEnd(inputStream, data, 0, data.length);
+
+    List<OggSeekPoint> seekPoints = new OggPageScanner(absoluteOffset, data, dataLength).createSeekTable(sampleRate);
+
+    inputStream.seek(savedPosition);
+    return seekPoints;
+  }
+
   /**
    * If it is possible to seek backwards on this stream, and the length of the stream is known, seeks to the end of the
    * track to determine the stream length both in bytes and samples.
@@ -310,8 +375,8 @@ private OggStreamSizeInfo scanForSizeInfo(int tailLength, int sampleRate) throws
     byte[] data = new byte[tailLength];
     int dataLength = StreamTools.readUntilEnd(inputStream, data, 0, data.length);
 
-    return new OggPageScanner(absoluteOffset, data, dataLength).scanForSizeInfo(pageHeader.byteStreamPosition,
-        sampleRate);
+    return new OggPageScanner(absoluteOffset, data, dataLength)
+        .scanForSizeInfo(pageHeader.byteStreamPosition, sampleRate);
   }
 
   /**
@@ -339,9 +404,10 @@ private boolean continuePacket() throws IOException {
   }
 
   private enum State {
+    TRACK_SEEKING,
     TRACK_BOUNDARY,
     PACKET_BOUNDARY,
     PACKET_READ,
     TERMINATED
   }
-}
+}

main/src/main/java/com/sedmelluq/discord/lavaplayer/container/ogg/OggPageScanner.java

@@ -1,6 +1,8 @@
 package com.sedmelluq.discord.lavaplayer.container.ogg;
 
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
 
 /**
  * Scanner for determining OGG stream information by seeking around in it.
@@ -16,6 +18,7 @@ public class OggPageScanner {
   private long reversedPosition;
   private int pageSize;
   private long byteStreamPosition;
+  private int pageSequence;
 
   /**
    * @param absoluteOffset Current position of the stream in bytes.
@@ -59,6 +62,38 @@ public OggStreamSizeInfo scanForSizeInfo(long firstPageOffset, int sampleRate) {
     return null;
   }
 
+  /**
+   * Creates a seek table for the OGG stream.
+   *
+   * @param sampleRate Sample rate of the track in the stream.
+   * @return A list of OggSeekPoint objects representing the seek points in the stream.
+   */
+  public List<OggSeekPoint> createSeekTable(int sampleRate) {
+    List<OggSeekPoint> seekPoints = new ArrayList<>();
+
+    ByteBuffer buffer = ByteBuffer.wrap(data, 0, dataLength);
+    int head = buffer.getInt(0);
+
+    for (int i = 0; i < dataLength - 27; i++) {
+      if (head == OGG_PAGE_HEADER_INT) {
+        buffer.position(i);
+
+        if (attemptReadHeader(buffer)) {
+          long position = byteStreamPosition;
+          long granulePosition = Long.reverseBytes(reversedPosition);
+          long timecode = granulePosition / (sampleRate / 1000);
+          pageSequence++;
+          seekPoints.add(new OggSeekPoint(position, granulePosition, timecode, pageSequence));
+        }
+      }
+
+      head <<= 8;
+      head |= data[i + 4] & 0xFF;
+    }
+
+    return seekPoints;
+  }
+
   private boolean attemptReadHeader(ByteBuffer buffer) {
     int start = buffer.position();
 

main/src/main/java/com/sedmelluq/discord/lavaplayer/container/ogg/OggSeekPoint.java

@@ -0,0 +1,49 @@
+package com.sedmelluq.discord.lavaplayer.container.ogg;
+
+public class OggSeekPoint {
+  private final long position;
+  private final long granulePosition;
+  private final long timecode;
+  private final long pageSequence;
+
+  /**
+   * @param position The position of the seek point in the stream, in bytes.
+   * @param granulePosition The granule position of the seek point in the stream.
+   * @param timecode The time of the seek point in the stream, in milliseconds.
+   * @param pageSequence The page to what this seek point belong.
+   */
+  public OggSeekPoint(long position, long granulePosition, long timecode, long pageSequence) {
+    this.position = position;
+    this.granulePosition = granulePosition;
+    this.timecode = timecode;
+    this.pageSequence = pageSequence;
+  }
+
+  /**
+   * @return The position of the seek point in the stream, in bytes.
+   */
+  public long getPosition() {
+    return position;
+  }
+
+  /**
+   * @return The granule position of the seek point in the stream.
+   */
+  public long getGranulePosition() {
+    return granulePosition;
+  }
+
+  /**
+   * @return The timecode of the seek point in the stream, in milliseconds.
+   */
+  public long getTimecode() {
+    return timecode;
+  }
+
+  /**
+   * @return The page to what this seek point belong.
+   */
+  public long getPageSequence() {
+    return pageSequence;
+  }
+}

main/src/main/java/com/sedmelluq/discord/lavaplayer/container/ogg/OggTrackBlueprint.java

@@ -2,4 +2,5 @@
 
 public interface OggTrackBlueprint {
   OggTrackHandler loadTrackHandler(OggPacketInputStream stream);
+  int getSampleRate();
 }

main/src/main/java/com/sedmelluq/discord/lavaplayer/container/ogg/flac/OggFlacCodecHandler.java

@@ -38,7 +38,9 @@ public int getMaximumFirstPacketLength() {
 
   @Override
   public OggTrackBlueprint loadBlueprint(OggPacketInputStream stream, DirectBufferStreamBroker broker) throws IOException {
-    return new Blueprint(load(stream, broker));
+    FlacTrackInfo info = load(stream, broker);
+    stream.setSeekPoints(stream.createSeekTable(info.stream.sampleRate));
+    return new Blueprint(info);
   }
 
   @Override
@@ -101,5 +103,10 @@ private Blueprint(FlacTrackInfo info) {
     public OggTrackHandler loadTrackHandler(OggPacketInputStream stream) {
       return new OggFlacTrackHandler(info, stream);
     }
+
+    @Override
+    public int getSampleRate() {
+      return info.stream.sampleRate;
+    }
   }
 }

main/src/main/java/com/sedmelluq/discord/lavaplayer/container/ogg/flac/OggFlacTrackHandler.java

@@ -71,7 +71,11 @@ private int readFlacFrame() throws IOException {
 
   @Override
   public void seekToTimecode(long timecode) {
-    throw new UnsupportedOperationException();
+    try {
+      downstream.seekPerformed(timecode, packetInputStream.seek(timecode));
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
   }
 
   @Override