diff --git a/mod/dagmodifier.go b/mod/dagmodifier.go index 9e86d7b60..83da608b9 100644 --- a/mod/dagmodifier.go +++ b/mod/dagmodifier.go @@ -152,23 +152,25 @@ func (dm *DagModifier) Write(b []byte) (int, error) { var ErrNoRawYet = fmt.Errorf("currently only fully support protonodes in the dagmodifier") func (dm *DagModifier) Size() (int64, error) { - pbnd, ok := dm.curNode.(*mdag.ProtoNode) - if !ok { - return 0, ErrNoRawYet - } - - pbn, err := ft.FromBytes(pbnd.Data()) - if err != nil { - return 0, err - } - - if dm.wrBuf != nil { - if uint64(dm.wrBuf.Len())+dm.writeStart > pbn.GetFilesize() { + switch nd := dm.curNode.(type) { + case *mdag.ProtoNode: + pbn, err := ft.FromBytes(nd.Data()) + if err != nil { + return 0, err + } + if dm.wrBuf != nil && uint64(dm.wrBuf.Len())+dm.writeStart > pbn.GetFilesize() { return int64(dm.wrBuf.Len()) + int64(dm.writeStart), nil } + return int64(pbn.GetFilesize()), nil + case *mdag.RawNode: + if dm.wrBuf != nil { + return 0, ErrNoRawYet + } + sz, err := nd.Size() + return int64(sz), err + default: + return 0, ErrNotUnixfs } - - return int64(pbn.GetFilesize()), nil } // Sync writes changes to this dag to disk @@ -397,12 +399,12 @@ func (dm *DagModifier) CtxReadFull(ctx context.Context, b []byte) (int, error) { } // GetNode gets the modified DAG Node -func (dm *DagModifier) GetNode() (*mdag.ProtoNode, error) { +func (dm *DagModifier) GetNode() (node.Node, error) { err := dm.Sync() if err != nil { return nil, err } - return dm.curNode.Copy().(*mdag.ProtoNode), nil + return dm.curNode.Copy(), nil } // HasChanges returned whether or not there are unflushed changes to this dag diff --git a/mod/dagmodifier_test.go b/mod/dagmodifier_test.go index d7b3f3267..b22844194 100644 --- a/mod/dagmodifier_test.go +++ b/mod/dagmodifier_test.go @@ -9,6 +9,7 @@ import ( h "github.com/ipfs/go-ipfs/importer/helpers" trickle "github.com/ipfs/go-ipfs/importer/trickle" + mdag "github.com/ipfs/go-ipfs/merkledag" ft "github.com/ipfs/go-ipfs/unixfs" uio "github.com/ipfs/go-ipfs/unixfs/io" testu "github.com/ipfs/go-ipfs/unixfs/test" @@ -105,7 +106,7 @@ func TestDagModifierBasic(t *testing.T) { t.Fatal(err) } - size, err := ft.DataSize(node.Data()) + size, err := ft.DataSize(node.(*mdag.ProtoNode).Data()) if err != nil { t.Fatal(err) }