module discovery

import attach

def refpat = /blake2#[0-9a-f]*/

test ManualDiscovery:
    let services = "discovery"

    subnet sd
    subnet s1
    subnet s2

    spawn as pd on sd
    spawn as p1 on s1
    spawn as p2 on s2
    send "create-identity Discovery" to pd
    send "create-identity Device1 Owner1" to p1
    send "create-identity Device2 Owner2" to p2

    expect /create-identity-done ref ($refpat).*/ from p1 capture p1id
    send "identity-info $p1id" to p1
    expect /identity-info ref $p1id base ($refpat) owner ($refpat).*/ from p1 capture p1base, p1owner
    send "identity-info $p1owner" to p1
    expect /identity-info ref $p1owner base ($refpat).*/ from p1 capture p1obase

    expect /create-identity-done ref $refpat.*/ from p2
    expect /create-identity-done ref $refpat.*/ from pd

    # Test discovery using owner and device identities:
    for id in [ p1obase, p1base ]:
        for p in [ pd, p1, p2 ]:
            send "start-server services $services" to p

        for p in [ p1, p2 ]:
            with p:
                send "peer-add ${pd.node.ip}"
                expect:
                    /peer 1 addr ${pd.node.ip} 29665/
                    /peer 1 id Discovery/
            expect from pd:
                    /peer [12] addr ${p.node.ip} 29665/
                    /peer [12] id .*/

        send "discovery-connect $id" to p2

        expect from p1:
            /peer [0-9]+ addr ${p2.node.ip} 29665/
            /peer [0-9]+ id Device2 Owner2/
        expect from p2:
            /peer [0-9]+ addr ${p1.node.ip} 29665/
            /peer [0-9]+ id Device1 Owner1/

        for p in [ pd, p1, p2 ]:
            send "stop-server" to p
        for p in [ pd, p1, p2 ]:
            expect /stop-server-done/ from p

    # Test delayed discovery with new peer
    for id in [ p1obase ]:
        for p in [ pd, p1, p2 ]:
            send "start-server services $services" to p

        with p1:
            send "peer-add ${pd.node.ip}"
            expect:
                /peer 1 addr ${pd.node.ip} 29665/
                /peer 1 id Discovery/
        expect from pd:
                /peer [12] addr ${p1.node.ip} 29665/
                /peer [12] id Device1 Owner1/

        send "discovery-connect $id" to p2

        with p2:
            send "peer-add ${pd.node.ip}"
            expect:
                /peer 1 addr ${pd.node.ip} 29665/
                /peer 1 id Discovery/
        expect from pd:
                /peer [12] addr ${p2.node.ip} 29665/
                /peer [12] id Device2 Owner2/

        expect from p1:
            /peer [0-9]+ addr ${p2.node.ip} 29665/
            /peer [0-9]+ id Device2 Owner2/
        expect from p2:
            /peer [0-9]+ addr ${p1.node.ip} 29665/
            /peer [0-9]+ id Device1 Owner1/

        for p in [ pd, p1, p2 ]:
            send "stop-server" to p
        for p in [ pd, p1, p2 ]:
            expect /stop-server-done/ from p


test DiscoveryJoinedAfter:
    let services = "discovery"

    subnet sd
    subnet s1
    subnet s2

    spawn as pd on sd
    spawn as p1 on s1
    spawn as p2 on s2
    send "create-identity Discovery" to pd
    send "create-identity Device1 Owner1" to p1
    send "create-identity Device2 Owner2" to p2

    expect /create-identity-done ref ($refpat).*/ from p1 capture p1id
    send "identity-info $p1id" to p1
    expect /identity-info ref $p1id base ($refpat) owner ($refpat).*/ from p1 capture p1base, p1owner
    send "identity-info $p1owner" to p1
    expect /identity-info ref $p1owner base ($refpat).*/ from p1 capture p1obase

    expect /create-identity-done ref $refpat.*/ from p2
    expect /create-identity-done ref $refpat.*/ from pd

    for id in [ p1obase ]:
        for p in [ pd, p1, p2 ]:
            send "start-server services $services" to p

        for p in [ p2 ]:
            with p:
                send "peer-add ${pd.node.ip}"
                expect:
                    /peer 1 addr ${pd.node.ip} 29665/
                    /peer 1 id Discovery/
            expect from pd:
                    /peer [12] addr ${p.node.ip} 29665/
                    /peer [12] id Device. Owner./

        send "discovery-connect $id" to p2

        for p in [ p1 ]:
            with p:
                send "peer-add ${pd.node.ip}"
                expect:
                    /peer 1 addr ${pd.node.ip} 29665/
                    /peer 1 id Discovery/
            expect from pd:
                    /peer [12] addr ${p.node.ip} 29665/
                    /peer [12] id Device. Owner./

        expect from p1:
            /peer [0-9]+ addr ${p2.node.ip} 29665/
            /peer [0-9]+ id Device2 Owner2/
        expect from p2:
            /peer [0-9]+ addr ${p1.node.ip} 29665/
            /peer [0-9]+ id Device1 Owner1/

        for p in [ pd, p1, p2 ]:
            send "stop-server" to p
        for p in [ pd, p1, p2 ]:
            expect /stop-server-done/ from p


test DiscoveryTunnel:
    let services = "discovery:tunnel"

    subnet sd
    subnet s1
    subnet s2

    spawn as pd on sd
    spawn as p1 on s1
    spawn as p2 on s2

    for n in [ p1.node, p2.node ]:
        shell on n:
            nft add table inet filter
            nft add chain inet filter input '{ type filter hook input priority filter ; policy drop; }'
            nft add rule inet filter input 'ct state { established, related } accept'

    send "create-identity Discovery" to pd
    send "create-identity Device1 Owner1" to p1
    send "create-identity Device2 Owner2" to p2

    expect /create-identity-done ref ($refpat).*/ from p1 capture p1id
    send "identity-info $p1id" to p1
    expect /identity-info ref $p1id base ($refpat) owner ($refpat).*/ from p1 capture p1base, p1owner
    send "identity-info $p1owner" to p1
    expect /identity-info ref $p1owner base ($refpat).*/ from p1 capture p1obase

    expect /create-identity-done ref $refpat.*/ from p2
    expect /create-identity-done ref $refpat.*/ from pd

    for id in [ p1obase ]:
        for p in [ pd, p1, p2 ]:
            send "start-server services $services test-log" to p

        for p in [ p1, p2 ]:
            with p:
                send "peer-add ${pd.node.ip}"
                expect:
                    /peer 1 addr ${pd.node.ip} 29665/
                    /peer 1 id Discovery/
            expect from pd:
                    /peer [12] addr ${p.node.ip} 29665/
                    /peer [12] id .*/

        send "discovery-tunnel 1 $id" to p2

        expect /net-ostream-open ${pd.node.ip} 29665 1 1/ from p2
        expect /net-ostream-open ${p1.node.ip} 29665 1 1/ from pd
        expect /net-ostream-open ${pd.node.ip} 29665 1 1/ from p1
        expect /net-ostream-open ${p2.node.ip} 29665 1 1/ from pd

        expect from p1:
            /peer 2 addr tunnel@.*/
            /peer 2 id Device2 Owner2/
        expect from p2:
            /peer 2 addr tunnel@.*/
            /peer 2 id Device1 Owner1/

        send "peer-drop 2" to p1
        send "peer-drop 2" to p2

        expect /net-ostream-close-ack ${pd.node.ip} 29665 1 0/ from p2
        expect /net-ostream-close-ack ${p1.node.ip} 29665 1 0/ from pd
        expect /net-ostream-close-ack ${pd.node.ip} 29665 1 0/ from p1
        expect /net-ostream-close-ack ${p2.node.ip} 29665 1 0/ from pd

        for p in [ pd, p1, p2 ]:
            send "stop-server" to p
        for p in [ pd, p1, p2 ]:
            expect /stop-server-done/ from p


test DiscoveryTunnelWithChangedAddresses:
    let services = "discovery:tunnel"

    subnet sd
    subnet s1
    subnet s2

    spawn as pd on sd
    spawn as p1 on s1
    spawn as p2 on s2

    for n in [ p1.node, p2.node ]:
        shell on n:
            nft add table inet filter
            nft add chain inet filter input '{ type filter hook input priority filter ; policy drop; }'
            nft add rule inet filter input 'ct state { established, related } accept'

    send "create-identity Discovery" to pd
    send "create-identity Device1 Owner1" to p1
    send "create-identity Device2 Owner2" to p2

    expect /create-identity-done ref ($refpat).*/ from p1 capture p1id
    send "identity-info $p1id" to p1
    expect /identity-info ref $p1id base ($refpat) owner ($refpat).*/ from p1 capture p1base, p1owner
    send "identity-info $p1owner" to p1
    expect /identity-info ref $p1owner base ($refpat).*/ from p1 capture p1obase

    expect /create-identity-done ref $refpat.*/ from p2
    expect /create-identity-done ref $refpat.*/ from pd

    for id in [ p1obase ]:
        send "start-server services $services test-log" to pd

        for p in [ p1, p2 ]:
            send "start-server services $services test-log" to p

        for p in [ p1, p2 ]:
            with p:
                send "peer-add ${pd.node.ip}"
                expect:
                    /peer 1 addr ${pd.node.ip} 29665/
                    /peer 1 id Discovery/
            expect from pd:
                    /peer [12] addr ${p.node.ip} 29665/
                    /peer [12] id .*/

        send "peer-drop 1" to pd
        send "peer-drop 2" to pd
        for p in [ p1, p2 ]:
            send "stop-server" to p
        for p in [ p1, p2 ]:
            shell on p.node as get_default:
                ip route | grep ^default
            expect /.* via ([0-9.]+).*/ from get_default capture default_route
            shell on p.node:
                ip addr del dev veth0 ${p.node.ip}/24
                ip addr add dev veth0 ${p.node.ip}0/24
                ip route add default via ${default_route} dev veth0 src ${p.node.ip}0
        for p in [ p1, p2 ]:
            expect /stop-server-done/ from p
        for p in [ p1, p2 ]:
            send "start-server services $services test-log" to p

        for p in [ p1, p2 ]:
            with p:
                send "peer-add ${pd.node.ip}"
                expect:
                    /peer 1 addr ${pd.node.ip} 29665/
                    /peer 1 id Discovery/
            expect from pd:
                    /peer [0-9]+ addr ${p.node.ip}0 29665/
                    /peer [0-9]+ id .*/

        send "discovery-tunnel 1 $id" to p2

        expect /net-ostream-open ${pd.node.ip} 29665 1 1/ from p2
        expect /net-ostream-open ${p1.node.ip}0 29665 1 1/ from pd
        expect /net-ostream-open ${pd.node.ip} 29665 1 1/ from p1
        expect /net-ostream-open ${p2.node.ip}0 29665 1 1/ from pd

        for p in [ pd, p1, p2 ]:
            send "stop-server" to p
        for p in [ pd, p1, p2 ]:
            expect /stop-server-done/ from p


test DiscoveryTunnelRefused:
    let services = "discovery"

    subnet sd
    subnet s1
    subnet s2

    spawn as pd on sd
    spawn as p1 on s1
    spawn as p2 on s2

    for n in [ p1.node, p2.node ]:
        shell on n:
            nft add table inet filter
            nft add chain inet filter input '{ type filter hook input priority filter ; policy drop; }'
            nft add rule inet filter input 'ct state { established, related } accept'

    send "create-identity Discovery" to pd
    send "create-identity Device1 Owner1" to p1
    send "create-identity Device2 Owner2" to p2

    expect /create-identity-done ref ($refpat).*/ from p1 capture p1id
    send "identity-info $p1id" to p1
    expect /identity-info ref $p1id base ($refpat) owner ($refpat).*/ from p1 capture p1base, p1owner
    send "identity-info $p1owner" to p1
    expect /identity-info ref $p1owner base ($refpat).*/ from p1 capture p1obase

    expect /create-identity-done ref $refpat.*/ from p2
    expect /create-identity-done ref $refpat.*/ from pd

    for id in [ p1obase ]:
        for p in [ pd, p1, p2 ]:
            send "start-server services $services test-log" to p

        for p in [ p1, p2 ]:
            with p:
                send "peer-add ${pd.node.ip}"
                expect:
                    /peer 1 addr ${pd.node.ip} 29665/
                    /peer 1 id Discovery/
            expect from pd:
                    /peer [12] addr ${p.node.ip} 29665/
                    /peer [12] id .*/

        send "discovery-tunnel 1 $id" to p2
        expect /net-ostream-open ${pd.node.ip} 29665 1 1/ from p2
        expect /net-ostream-close-ack ${pd.node.ip} 29665 1 0/ from p2

        for p in [ pd, p1, p2 ]:
            send "stop-server" to p
        for p in [ pd, p1, p2 ]:
            expect /stop-server-done/ from p


test CommonOwnerDiscovery:
    let services = "discovery:trace"

    subnet sd
    subnet s1
    subnet s2
    subnet s3

    node nd on sd
    node n1 on s1
    node n2 on s2
    node n3 on s3

    # TODO: move under the following "local" and capture values from inner block
    spawn as p1_init on n1
    spawn as p2_init on n2
    spawn as p3_init on n3

    send "create-identity Device1 Owner" to p1_init
    expect /create-identity-done ref ($refpat).*/ from p1_init capture p1_id_ext
    send "identity-info $p1_id_ext" to p1_init
    expect /identity-info ref $p1_id_ext base ($refpat) owner ($refpat).*/ from p1_init capture p1_id, owner_ext
    send "identity-info $owner_ext" to p1_init
    expect /identity-info ref $owner_ext base ($refpat).*/ from p1_init capture owner

    local:
        spawn as pd on nd
        let p1 = p1_init
        let p2 = p2_init
        let p3 = p3_init

        send "create-identity Device2" to p2
        send "create-identity Device3" to p3

        expect /create-identity-done .*/ from p2
        expect /create-identity-done .*/ from p3

        send "create-identity Discovery" to pd

        expect /create-identity-done .*/ from pd

        for p in [ p1, p2, p3 ]:
            send "start-server services attach" to p
        for p in [ p2, p3 ]:
            send "watch-local-identity" to p
            send "peer-add ${p1.node.ip} 29665" to p
            expect /peer [0-9]+ addr ${p1.node.ip} 29665/ from p
            expect /peer [0-9]+ id .*/ from p
            perform_attach of p to p1

            expect /local-identity ($refpat) Device. .*/ from p capture p_id_ext
            send "identity-info $p_id_ext" to p

    for p in [ p1_init, p2_init, p3_init ]:
        send "stop-server" to p
        expect /stop-server-done/ from p
    expect /identity-info ref $refpat base ($refpat).*/ from p2_init capture p2_id
    expect /identity-info ref $refpat base ($refpat).*/ from p3_init capture p3_id

    local:
        spawn as pd on nd
        spawn as p1 on n1
        spawn as p2 on n2
        spawn as p3 on n3

        ignore from p1 matching /discovery-packet-result via $refpat addr ICE/
        ignore from p2 matching /discovery-packet-result via $refpat addr ICE/
        ignore from p3 matching /discovery-packet-result via $refpat addr ICE/

        # TODO: direct result should be returned only for exact device id match, keep for now
        ignore from p1 matching /discovery-packet-result result .*/
        ignore from p2 matching /discovery-packet-result result .*/
        ignore from p3 matching /discovery-packet-result result .*/

        for p in [ pd, p1, p2, p3 ]:
            send "start-server services $services" to p

        send "peer-add ${pd.node.ip}" to p1
        expect from pd /discovery-packet ${p1.node.ip} 29665 search $owner/

        send "peer-add ${pd.node.ip}" to p2
        expect from pd /discovery-packet ${p2.node.ip} 29665 search $owner/

        expect from p1 /discovery-packet ${pd.node.ip} 29665 result $owner/
        expect from p1 /discovery-packet-result via ${p2_id} addr ${p2.node.ip} 29665/
        local:
            expect from p1 /discovery-packet-result (.*)/ capture done
            guard (done == done)

        expect from p2 /discovery-packet ${pd.node.ip} 29665 result $owner/
        expect from p2 /discovery-packet-result via ${p1_id} addr ${p1.node.ip} 29665/
        local:
            expect from p2 /discovery-packet-result (.*)/ capture done
            guard (done == "done")


        send "peer-add ${pd.node.ip}" to p3
        expect from pd /discovery-packet ${p3.node.ip} 29665 search $owner/

        expect from p3 /discovery-packet ${pd.node.ip} 29665 result $owner/
        expect from p3 /discovery-packet-result via ${p1_id} addr ${p1.node.ip} 29665/
        expect from p3 /discovery-packet-result via ${p2_id} addr ${p2.node.ip} 29665/
        local:
            expect from p3 /discovery-packet-result (.*)/ capture done
            guard (done == done)

        expect from p1:
            /peer [0-9]+ addr ${p2.node.ip} 29665/
            /peer [0-9]+ addr ${p3.node.ip} 29665/