contact.go

package kademlia

import (
  "fmt";
)

type Contact struct {
  id NodeID;
  address string;
}

func (contact *Contact) String() string {
  return fmt.Sprintf("Contact(\"%s\", \"%s\")", contact.id, contact.address);
}

kademlia.go

package kademlia

import (
  "container/heap";
  "container/vector";
  "fmt";
  "log";
  "http";
  "net";
  "os";
  "rpc";
  "sort";
)

type Kademlia struct {
  routes *RoutingTable;
  NetworkId string;
}

func NewKademlia(self *Contact, networkId string) (ret *Kademlia) {
  ret = new(Kademlia);
  ret.routes = NewRoutingTable(self);
  ret.NetworkId = networkId;
  return;
}

func (k *Kademlia) Serve() (err os.Error) {
  rpc.Register(&KademliaCore{k});

  rpc.HandleHTTP();
  if l, err := net.Listen("tcp", k.routes.node.address); err == nil {
    go http.Serve(l, nil);
  }
  return;
}

func (k *Kademlia) Call(contact *Contact, method string, args, reply interface{}) (err os.Error) {
  if client, err := rpc.DialHTTP("tcp", contact.address); err == nil {
    err = client.Call(method, args, reply);
    if err == nil {
      k.routes.Update(contact);
    }
  }
  return;
}

func (k *Kademlia) sendQuery(node *Contact, target NodeID, done chan []Contact) {
  args := FindNodeRequest{RPCHeader{&k.routes.node, k.NetworkId}, target};
  reply := FindNodeResponse{};
  
  if err := k.Call(node, "KademliaCore.FindNode", &args, &reply); err == nil {
    done <- reply.contacts;
  } else {
    done <- []Contact{};
  }
}

func (k *Kademlia) IterativeFindNode(target NodeID, delta int) (ret *vector.Vector) {
  done := make(chan []Contact);

  // A vector of *ContactRecord structs
  ret = new(vector.Vector).Resize(0, BucketSize);
  
  // A heap of not-yet-queried *Contact structs
  frontier := new(vector.Vector).Resize(0, BucketSize);
  
  // A map of client values we've seen so far
  seen := make(map[string] bool);
  
  // Initialize the return list, frontier heap, and seen list with local nodes
  for node := range k.routes.FindClosest(target, delta).Iter() {
    record := node.(*ContactRecord);
    ret.Push(record);
    heap.Push(frontier, record.node);
    seen[record.node.id.String()] = true;
  }
  
  // Start off delta queries
  pending := 0;
  for i := 0; i < delta && frontier.Len() > 0; i++ {
    pending++;
    go k.sendQuery(frontier.Pop().(*Contact), target, done);
  }
  
  // Iteratively look for closer nodes
  for pending > 0 {
    nodes := <-done;
    pending--;
    for _, node := range nodes {
      // If we haven't seen the node before, add it
      if _, ok := seen[node.id.String()]; ok == false {
        ret.Push(&ContactRecord{&node, node.id.Xor(target)});
        heap.Push(frontier, node);
        seen[node.id.String()] = true;
      }
    }
    
    for pending < delta && frontier.Len() > 0 {
      go k.sendQuery(frontier.Pop().(*Contact), target, done);
      pending++;
    }
  }

  sort.Sort(ret);
  if ret.Len() > BucketSize {
    ret.Cut(BucketSize, ret.Len());
  }

  return;
}

type RPCHeader struct {
  Sender *Contact;
  NetworkId string;
}

func (k *Kademlia) HandleRPC(request, response *RPCHeader) os.Error {
  if request.NetworkId != k.NetworkId {
    return os.NewError(fmt.Sprintf("Expected network ID %s, got %s",
                                   k.NetworkId, request.NetworkId));
  }
  if request.Sender != nil {
    k.routes.Update(request.Sender);
  }
  response.Sender = &k.routes.node;
  return nil;
}

type KademliaCore struct {
  kad *Kademlia;
}

type PingRequest struct {
  RPCHeader;
}

type PingResponse struct {
  RPCHeader;
}

func (kc *KademliaCore) Ping(args *PingRequest, response *PingResponse) (err os.Error) {
  if err = kc.kad.HandleRPC(&args.RPCHeader, &response.RPCHeader); err == nil {
    log.Stderr("Ping from %s\n", args.RPCHeader);
  }
  return;
}

type FindNodeRequest struct {
  RPCHeader;
  target NodeID;
}

type FindNodeResponse struct {
  RPCHeader;
  contacts []Contact;
}

func (kc *KademliaCore) FindNode(args *FindNodeRequest, response *FindNodeResponse) (err os.Error) {
  if err = kc.kad.HandleRPC(&args.RPCHeader, &response.RPCHeader); err == nil {
    contacts := kc.kad.routes.FindClosest(args.target, BucketSize);
    response.contacts = make([]Contact, contacts.Len());

    for i := 0; i < contacts.Len(); i++ {
      response.contacts[i] = *contacts.At(i).(*ContactRecord).node;
    }
  }
  return;
}

kademlia_test.go

package kademlia

import "testing"

func TestPing(t *testing.T) {
  me := Contact{NewRandomNodeID(), "127.0.0.1:8989"};
  k := NewKademlia(&me, "test");
  k.Serve();

  someone := Contact{NewRandomNodeID(), "127.0.0.1:8989"};
  if err := k.Call(
      &someone,
      "KademliaCore.Ping",
      &PingRequest{RPCHeader{&someone, k.NetworkId}},
      &PingResponse{}); err != nil {
    t.Error(err);
  }
}

func TestFindNode(t *testing.T) {
  me := Contact{NewRandomNodeID(), "127.0.0.1:8989"};
  k := NewKademlia(&me, "test");
  kc := KademliaCore{k};
  
  var contacts [100]Contact;
  for i := 0; i < len(contacts); i++ {
    contacts[i] = Contact{NewRandomNodeID(), "127.0.0.1:8989"};
    if err := kc.Ping(&PingRequest{RPCHeader{&contacts[i], k.NetworkId}},
                      &PingResponse{}); err != nil {
      t.Error(err);
    }
  }

  args := FindNodeRequest{RPCHeader{&contacts[0], k.NetworkId}, contacts[0].id};
  response := FindNodeResponse{};
  if err := kc.FindNode(&args, &response); err != nil {
    t.Error(err);
  }
  
  if len(response.contacts) != BucketSize {
    t.Fail();
  }
}

Makefile

include $(GOROOT)/src/Make.$(GOARCH)

TARG=kademlia
GOFILES=\
        nodeid.go\
        routingtable.go\
				contact.go\
				kademlia.go

include $(GOROOT)/src/Make.pkg

nodeid.go

package kademlia

import (
  "encoding/hex";
  "rand";
)

const IdLength = 20

type NodeID [IdLength]byte

func NewNodeID(data string) (ret NodeID) {
  decoded, _ := hex.DecodeString(data);
  for i := 0; i < IdLength; i++ {
    ret[i] = decoded[i];
  }
  return;
}

func NewRandomNodeID() (ret NodeID) {
  for i := 0; i < IdLength; i++ {
    ret[i] = uint8(rand.Intn(256));
  }
  return;
}

func (node NodeID) String() string {
  return hex.EncodeToString(node[0:IdLength]);
}

func (node NodeID) Equals(other NodeID) bool {
  for i := 0; i < IdLength; i++ {
    if node[i] != other[i] {
      return false;
    }
  }
  return true;
}

func (node NodeID) Xor(other NodeID) (ret NodeID) {
  for i := 0; i < IdLength; i++ {
    ret[i] = node[i] ^ other[i];
  }
  return;
}

func (node NodeID) PrefixLen() (ret int) {
  for i := 0; i < IdLength; i++ {
    for j := 0; j < 8; j++ {
      if (node[i] >> uint8(7 - j)) & 0x1 != 0 {
        return i * 8 + j;
      }
    }
  }
  return IdLength * 8 - 1;
}

func (node NodeID) Less(other interface{}) bool {
  for i := 0; i < IdLength; i++ {
    if node[i] != other.(NodeID)[i] {
      return node[i] < other.(NodeID)[i];
    }
  }
  return false;
}

nodeid_test.go

package kademlia

import (
  "os";
  "testing";
  "fmt";
)

func TestNodeID(t *testing.T) {
  a := NodeID{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19};
  b := NodeID{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 19, 18};
  c := NodeID{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,  0,  0,  0,  0,  0,  0,  0,  1,  1};

  if !a.Equals(a) {
    t.Fail();
  }
  if a.Equals(b) {
    t.Fail();
  }
  
  if !a.Xor(b).Equals(c) {
    t.Error(a.Xor(b));
  }
  
  if c.PrefixLen() != 151 {
    t.Error(c.PrefixLen());
  }
  
  if b.Less(a) {
    t.Fail();
  }
  
  str_id := "0123456789abcdef0123456789abcdef01234567";
  if NewNodeID(str_id).String() != str_id {
    t.Error(NewNodeID(str_id).String());
  }
}

type Foo struct {
  a int;
}

func (f Foo) String() string {
  return "foo";
}

type FmtWriter struct {
}

func (fw FmtWriter) Write(p []byte) (n int, err os.Error) {
  fmt.Printf("gob(%d): %x\n", len(p), p);
  return len(p), nil;
}

routingtable.go

package kademlia

import (
  "container/list";
  "container/vector";
  "exp/iterable";
  "sort";
)

const BucketSize = 20;

type RoutingTable struct {
  node Contact;
  buckets [IdLength*8]*list.List;
}

type ContactRecord struct {
  node *Contact;
  sortKey NodeID;
}

func (rec *ContactRecord) Less(other interface{}) bool {
  return rec.sortKey.Less(other.(*ContactRecord).sortKey);
}

func NewRoutingTable(node *Contact) (ret *RoutingTable) {
  ret = new(RoutingTable);
  for i := 0; i < IdLength * 8; i++ {
    ret.buckets[i] = list.New();
  }
  ret.node = *node;
  return;
}

func (table *RoutingTable) Update(contact *Contact) {
  prefix_length := contact.id.Xor(table.node.id).PrefixLen();
  bucket := table.buckets[prefix_length];
  element := iterable.Find(bucket, func(x interface{}) bool {
    return x.(*Contact).id.Equals(table.node.id);
  });
  if element == nil {
    if bucket.Len() <= BucketSize {
      bucket.PushFront(contact);
    }
    // TODO: Handle insertion when the list is full by evicting old elements if
    // they don't respond to a ping.
  } else {
    bucket.MoveToFront(element.(*list.Element));
  }
}

func copyToVector(start, end *list.Element, vec *vector.Vector, target NodeID) {
  for elt := start; elt != end; elt = elt.Next() {
    contact := elt.Value.(*Contact);
    vec.Push(&ContactRecord{contact, contact.id.Xor(target)});
  }
}

func (table *RoutingTable) FindClosest(target NodeID, count int) (ret *vector.Vector) {
  ret = new(vector.Vector).Resize(0, count);

  bucket_num := target.Xor(table.node.id).PrefixLen();
  bucket := table.buckets[bucket_num];
  copyToVector(bucket.Front(), nil, ret, target);

  for i := 1; (bucket_num-i >= 0 || bucket_num+i < IdLength*8) && ret.Len() < count; i++ {
    if bucket_num - i >= 0 {
      bucket = table.buckets[bucket_num - i];
      copyToVector(bucket.Front(), nil, ret, target);
    }
    if bucket_num + i < IdLength * 8 {
      bucket = table.buckets[bucket_num + i];
      copyToVector(bucket.Front(), nil, ret, target);
    }
  }
  
  sort.Sort(ret);
  if ret.Len() > count {
    ret.Cut(count, ret.Len());
  }
  return;
}

routingtable_test.go

package kademlia

import "testing"

func TestRoutingTable(t *testing.T) {
  n1 := NewNodeID("FFFFFFFF00000000000000000000000000000000");
  n2 := NewNodeID("FFFFFFF000000000000000000000000000000000");
  n3 := NewNodeID("1111111100000000000000000000000000000000");
  rt := NewRoutingTable(&Contact{n1, "localhost:8000"});
  rt.Update(&Contact{n2, "localhost:8001"});
  rt.Update(&Contact{n3, "localhost:8002"});
  
  vec := rt.FindClosest(NewNodeID("2222222200000000000000000000000000000000"), 1);
  if vec.Len() != 1 {
    t.Fail();
    return;
  }
  if !vec.At(0).(*ContactRecord).node.id.Equals(n3) {
    t.Error(vec.At(0));
  }
  
  vec = rt.FindClosest(n2, 10);
  if vec.Len() != 2 {
    t.Error(vec.Len());
    return;
  }
  if !vec.At(0).(*ContactRecord).node.id.Equals(n2) {
    t.Error(vec.At(0));
  }
  if !vec.At(1).(*ContactRecord).node.id.Equals(n3) {
    t.Error(vec.At(1));
  }
}