github.com/omniscale/go-osm@v0.3.1/parser/pbf/lowlevel.go (about)

     1  package pbf
     2  
     3  import (
     4  	"bytes"
     5  	"compress/zlib"
     6  	structs "encoding/binary"
     7  	"errors"
     8  	"fmt"
     9  	"io"
    10  	"time"
    11  
    12  	"github.com/gogo/protobuf/proto"
    13  	"github.com/omniscale/go-osm/parser/pbf/internal/osmpbf"
    14  )
    15  
    16  var supportedFeatured = map[string]bool{"OsmSchema-V0.6": true, "DenseNodes": true}
    17  
    18  // decodeRawBlob decodes Blob PBF messages and returns either the raw bytes or
    19  // the uncompressed zlib_data bytes. The result can contain encoded HeaderBlock
    20  // or PrimitiveBlock PBF messages.
    21  func decodeRawBlob(raw []byte) ([]byte, error) {
    22  	blob := &osmpbf.Blob{}
    23  
    24  	err := proto.Unmarshal(raw, blob)
    25  	if err != nil {
    26  		return nil, fmt.Errorf("unmarshaling blob: %w", err)
    27  	}
    28  
    29  	// pbf contains (uncompressed) raw or zlibdata
    30  	b := blob.GetRaw()
    31  	if b == nil {
    32  		buf := bytes.NewBuffer(blob.GetZlibData())
    33  		r, err := zlib.NewReader(buf)
    34  		if err != nil {
    35  			return nil, fmt.Errorf("start uncompressing ZLibData: %w", err)
    36  		}
    37  		b = make([]byte, blob.GetRawSize())
    38  		_, err = io.ReadFull(r, b)
    39  		if err != nil {
    40  			return nil, fmt.Errorf("uncompressing ZLibData: %w", err)
    41  		}
    42  	}
    43  	return b, nil
    44  }
    45  
    46  func decodePrimitiveBlock(blob []byte) (*osmpbf.PrimitiveBlock, error) {
    47  	b, err := decodeRawBlob(blob)
    48  	if err != nil {
    49  		return nil, fmt.Errorf("decoding raw blob: %w", err)
    50  	}
    51  	block := &osmpbf.PrimitiveBlock{}
    52  	if err = proto.Unmarshal(b, block); err != nil {
    53  		return nil, fmt.Errorf("unmarshaling PrimitiveBlock: %w", err)
    54  	}
    55  	return block, nil
    56  }
    57  
    58  func decodeHeaderBlock(blob []byte) (*Header, error) {
    59  	b, err := decodeRawBlob(blob)
    60  	if err != nil {
    61  		return nil, err
    62  	}
    63  
    64  	header := &osmpbf.HeaderBlock{}
    65  	if err := proto.Unmarshal(b, header); err != nil {
    66  		return nil, err
    67  	}
    68  
    69  	for _, feature := range header.RequiredFeatures {
    70  		if supportedFeatured[feature] != true {
    71  			return nil, fmt.Errorf("cannot parse file, feature %v not supported: %w", feature, err)
    72  		}
    73  	}
    74  
    75  	result := &Header{}
    76  	timestamp := header.GetOsmosisReplicationTimestamp()
    77  	if timestamp != 0 {
    78  		// keep result.Time zero if timestamp is 0
    79  		result.Time = time.Unix(timestamp, 0)
    80  	}
    81  	result.Sequence = header.GetOsmosisReplicationSequenceNumber()
    82  	result.RequiredFeatures = header.RequiredFeatures
    83  	result.OptionalFeatures = header.OptionalFeatures
    84  	return result, nil
    85  }
    86  
    87  type Header struct {
    88  	Time     time.Time
    89  	Sequence int64
    90  
    91  	RequiredFeatures []string
    92  	OptionalFeatures []string
    93  }
    94  
    95  func parseHeader(r io.Reader) (*Header, error) {
    96  	blockHeader, data, err := nextBlock(r)
    97  	if err != nil {
    98  		return nil, fmt.Errorf("reading header: %w", err)
    99  	}
   100  	if blockHeader.GetType() != "OSMHeader" {
   101  		return nil, errors.New("invalid block type, expected OSMHeader, got " + blockHeader.GetType())
   102  	}
   103  	header, err := decodeHeaderBlock(data)
   104  	return header, err
   105  }
   106  
   107  func nextBlock(r io.Reader) (*osmpbf.BlobHeader, []byte, error) {
   108  	header, err := nextBlobHeader(r)
   109  	if err == io.EOF {
   110  		return nil, nil, err
   111  	}
   112  	if err != nil {
   113  		return nil, nil, fmt.Errorf("reading next block header: %w", err)
   114  	}
   115  	size := header.GetDatasize()
   116  
   117  	data := make([]byte, size)
   118  	n, err := io.ReadFull(r, data)
   119  	if err != nil {
   120  		return nil, nil, fmt.Errorf("reading next block: %w", err)
   121  	}
   122  	if n != int(size) {
   123  		return nil, nil, fmt.Errorf("reading next block, only got %d bytes instead of %d", n, size)
   124  	}
   125  	return header, data, nil
   126  }
   127  
   128  func nextBlobHeader(r io.Reader) (*osmpbf.BlobHeader, error) {
   129  	var size int32
   130  	err := structs.Read(r, structs.BigEndian, &size)
   131  	if err == io.EOF {
   132  		return nil, err
   133  	}
   134  	if err != nil {
   135  		return nil, fmt.Errorf("reading header size: %w", err)
   136  	}
   137  
   138  	var blobHeader = &osmpbf.BlobHeader{}
   139  
   140  	data := make([]byte, size)
   141  	n, err := io.ReadFull(r, data)
   142  	if err != nil {
   143  		return nil, fmt.Errorf("reading blob header: %w", err)
   144  	}
   145  	if n != int(size) {
   146  		return nil, fmt.Errorf("reading blob header, only got %d bytes instead of %d", n, size)
   147  	}
   148  
   149  	err = proto.Unmarshal(data, blobHeader)
   150  	if err != nil {
   151  		return nil, fmt.Errorf("unmarshaling header: %w", err)
   152  	}
   153  
   154  	return blobHeader, nil
   155  }