Adding a timeout to fix infinite looping on corrupt bdb database files.

Author: erikvarga

pkg/bdb/bdb.go

@@ -1,6 +1,7 @@
 package bdb
 
 import (
+	"context"
 	"io"
 	"os"
 
@@ -61,7 +62,7 @@ func (db *BerkeleyDB) Close() error {
 	return db.file.Close()
 }
 
-func (db *BerkeleyDB) Read() <-chan dbi.Entry {
+func (db *BerkeleyDB) Read(ctx context.Context) <-chan dbi.Entry {
 	entries := make(chan dbi.Entry)
 
 	go func() {
@@ -118,6 +119,7 @@ func (db *BerkeleyDB) Read() <-chan dbi.Entry {
 
 				// Traverse the page to concatenate the data that may span multiple pages.
 				valueContent, err := HashPageValueContent(
+					ctx,
 					db.file,
 					pageData,
 					hashPageIndex,

pkg/bdb/hash_page.go

@@ -2,6 +2,7 @@ package bdb
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"io"
 	"os"
@@ -31,7 +32,7 @@ func ParseHashPage(data []byte, swapped bool) (*HashPage, error) {
 	return &hashPage, nil
 }
 
-func HashPageValueContent(db *os.File, pageData []byte, hashPageIndex uint16, pageSize uint32, swapped bool) ([]byte, error) {
+func HashPageValueContent(ctx context.Context, db *os.File, pageData []byte, hashPageIndex uint16, pageSize uint32, swapped bool) ([]byte, error) {
 	// the first byte is the page type, so we can peek at it first before parsing further...
 	valuePageType := pageData[hashPageIndex]
 
@@ -50,6 +51,12 @@ func HashPageValueContent(db *os.File, pageData []byte, hashPageIndex uint16, pa
 	var hashValue []byte
 
 	for currentPageNo := entry.PageNo; currentPageNo != 0; {
+		select {
+		case <-ctx.Done():
+			return nil, xerrors.Errorf("timed out parsing hash page")
+		default:
+		}
+
 		pageStart := pageSize * currentPageNo
 
 		_, err := db.Seek(int64(pageStart), io.SeekStart)

pkg/db/rpmdbinterface.go

@@ -1,11 +1,13 @@
 package dbi
 
+import "context"
+
 type Entry struct {
 	Value []byte
 	Err   error
 }
 
 type RpmDBInterface interface {
-	Read() <-chan Entry
+	Read(ctx context.Context) <-chan Entry
 	Close() error
 }

pkg/ndb/ndb.go

@@ -22,6 +22,7 @@ SOFTWARE.
 package ndb
 
 import (
+	"context"
 	"encoding/binary"
 	"fmt"
 	"io"
@@ -136,7 +137,7 @@ func (db *RpmNDB) Close() error {
 	return db.file.Close()
 }
 
-func (db *RpmNDB) Read() <-chan dbi.Entry {
+func (db *RpmNDB) Read(ctx context.Context) <-chan dbi.Entry {
 	entries := make(chan dbi.Entry)
 
 	go func() {

pkg/rpmdb.go

@@ -1,6 +1,8 @@
 package rpmdb
 
 import (
+	"context"
+
 	"github.com/knqyf263/go-rpmdb/pkg/bdb"
 	dbi "github.com/knqyf263/go-rpmdb/pkg/db"
 	"github.com/knqyf263/go-rpmdb/pkg/ndb"
@@ -47,7 +49,11 @@ func (d *RpmDB) Close() error {
 }
 
 func (d *RpmDB) Package(name string) (*PackageInfo, error) {
-	pkgs, err := d.ListPackages()
+	return d.PackageWithContext(context.TODO(), name)
+}
+
+func (d *RpmDB) PackageWithContext(ctx context.Context, name string) (*PackageInfo, error) {
+	pkgs, err := d.ListPackagesWithContext(ctx)
 	if err != nil {
 		return nil, xerrors.Errorf("unable to list packages: %w", err)
 	}
@@ -61,9 +67,13 @@ func (d *RpmDB) Package(name string) (*PackageInfo, error) {
 }
 
 func (d *RpmDB) ListPackages() ([]*PackageInfo, error) {
+	return d.ListPackagesWithContext(context.TODO())
+}
+
+func (d *RpmDB) ListPackagesWithContext(ctx context.Context) ([]*PackageInfo, error) {
 	var pkgList []*PackageInfo
 
-	for entry := range d.db.Read() {
+	for entry := range d.db.Read(ctx) {
 		if entry.Err != nil {
 			return nil, entry.Err
 		}

pkg/rpmdb_test.go

@@ -1,8 +1,10 @@
 package rpmdb
 
 import (
+	"context"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -824,3 +826,16 @@ func TestNevra(t *testing.T) {
 	_, err = pkg.InstalledFiles()
 	require.Error(t, err)
 }
+
+func TestTimeoutPackages(t *testing.T) {
+	db, err := Open("testdata/centos7-many/Packages")
+	require.NoError(t, err)
+	ctxTimesOut, cancelFunc := context.WithTimeout(context.Background(), 1*time.Microsecond)
+	defer cancelFunc()
+	_, err = db.ListPackagesWithContext(ctxTimesOut)
+	if err == nil {
+		t.Errorf("Expected timeout parsing hash page")
+	} else {
+		assert.Equal(t, "timed out parsing hash page", err.Error())
+	}
+}

pkg/sqlite3/sqlite3.go

@@ -2,6 +2,7 @@ package sqlite3
 
 import (
 	"bytes"
+	"context"
 	"database/sql"
 	"encoding/binary"
 	"os"
@@ -44,7 +45,7 @@ func Open(path string) (*SQLite3, error) {
 	return &SQLite3{db}, nil
 }
 
-func (db *SQLite3) Read() <-chan dbi.Entry {
+func (db *SQLite3) Read(ctx context.Context) <-chan dbi.Entry {
 	entries := make(chan dbi.Entry)
 
 	go func() {