package peergos.shared.storage;

import peergos.shared.crypto.hash.*;
import peergos.shared.io.ipfs.cid.*;
import peergos.shared.io.ipfs.multihash.*;
import peergos.shared.util.*;

import java.util.*;
import java.util.concurrent.*;
import java.util.function.*;

public class WriteFilter extends DelegatingStorage {

    private final ContentAddressedStorage dht;
    private final BiFunction<PublicKeyHash, Integer, Boolean> keyFilter;

    public WriteFilter(ContentAddressedStorage dht, BiFunction<PublicKeyHash, Integer, Boolean> keyFilter) {
        super(dht);
        this.dht = dht;
        this.keyFilter = keyFilter;
    }

    @Override
    public CompletableFuture<List<PresignedUrl>> authWrites(PublicKeyHash owner,
                                                            PublicKeyHash writer,
                                                            List<byte[]> signedHashes,
                                                            List<Integer> blockSizes,
                                                            boolean isRaw,
                                                            TransactionId tid) {
        if (! keyFilter.apply(writer, blockSizes.stream().mapToInt(x -> x).sum()))
            throw new IllegalStateException("Key not allowed to write to this server: " + writer);
        return dht.authWrites(owner, writer, signedHashes, blockSizes, isRaw, tid);
    }

    @Override
    public ContentAddressedStorage directToOrigin() {
        return this;
    }

    @Override
    public CompletableFuture<List<Cid>> put(PublicKeyHash owner,
                                            PublicKeyHash writer,
                                            List<byte[]> signedHashes,
                                            List<byte[]> blocks,
                                            TransactionId tid) {
        if (! keyFilter.apply(writer, blocks.stream().mapToInt(x -> x.length).sum()))
            throw new IllegalStateException("Key not allowed to write to this server: " + writer);
        return dht.put(owner, writer, signedHashes, blocks, tid);
    }

    @Override
    public CompletableFuture<List<Cid>> putRaw(PublicKeyHash owner,
                                               PublicKeyHash writer,
                                               List<byte[]> signatures,
                                               List<byte[]> blocks,
                                               TransactionId tid,
                                               ProgressConsumer<Long> progressConsumer) {
        if (! keyFilter.apply(writer, blocks.stream().mapToInt(x -> x.length).sum()))
            throw new IllegalStateException("Key not allowed to write to this server: " + writer);
        return dht.putRaw(owner, writer, signatures, blocks, tid, progressConsumer);
    }
}

@@ -126,25 +126,15 @@
 red.
-io.ipfs.multihash
+user.fs
 .*;%0A
@@ -1296,32 +1296,183 @@
 r: %22 + writer);%0A
+        if (blockSizes.stream().anyMatch(s -%3E s %3E Fragment.MAX_LENGTH_WITH_BAT_PREFIX))%0A            throw new IllegalStateException(%22Block too big!%22);%0A
         return d
@@ -2165,32 +2165,186 @@
 r: %22 + writer);%0A
+        if (blocks.stream().anyMatch(b -%3E b.length %3E Fragment.MAX_LENGTH_WITH_BAT_PREFIX))%0A            throw new IllegalStateException(%22Block too big!%22);%0A
         return d
@@ -2925,32 +2925,32 @@
 length).sum()))%0A
-
             thro
@@ -3024,32 +3024,186 @@
 r: %22 + writer);%0A
+        if (blocks.stream().anyMatch(b -%3E b.length %3E Fragment.MAX_LENGTH_WITH_BAT_PREFIX))%0A            throw new IllegalStateException(%22Block too big!%22);%0A
         return d

Put in extra guards against too large blocks