Source

gotftp / gotftp.go

package main

import (
	"net"
	"io/ioutil"
	"os"
	"log"
	"strings"
)

// TFTP error codes
const (
	RRQ  = 1
	WRQ  = 2
	DATA = 3
	ACK  = 4
	ERR  = 5
)

// TFTP modes
const (
	MODE_ASCII = "netascii"
	MODE_OCTET = "octet"
)

// Default TFTP chunk size
const CHUNK_SIZE = 512

// Single TFTP file data chunk
type TftpChunk struct {
	id   int
	data []byte
	ack  bool
}

// Converts data chunk to TFTP data packet
func (c *TftpChunk) Packet() (p []byte) {
	p = make([]byte, len(c.data)+4)
	p[0] = 0
	p[1] = DATA
	p[2] = uint8(c.id / 256)
	p[3] = uint8(c.id % 256)
	copy(p[4:], c.data[:])
	log.Println("packet id =", c.id, "len =", (len(p) - 4))
	return
}

// TFTP file represented as array of chunks
type TftpFile struct {
	chunks []TftpChunk
}

// Create new TFTP file object from local file
func NewTftpFile(name string) (f *TftpFile, err os.Error) {
	var data []byte

	if data, err = ioutil.ReadFile(name); err != nil {
		return
	}

	emptyChunk := (len(data)%CHUNK_SIZE == 0)
	numChunks := len(data)/CHUNK_SIZE + 1

	f = new(TftpFile)
	if emptyChunk {
		f.chunks = make([]TftpChunk, numChunks+1)
		f.chunks[numChunks].ack = false
		f.chunks[numChunks].data = nil
		f.chunks[numChunks].id = numChunks + 1
	} else {
		f.chunks = make([]TftpChunk, numChunks)
	}

	for i := 0; i < numChunks; i++ {
		//f.chunks[i] = new(TftpChunk)
		f.chunks[i].ack = false
		if i == numChunks-1 {
			f.chunks[i].data = data[i*CHUNK_SIZE : len(data)]
		} else {
			f.chunks[i].data = data[i*CHUNK_SIZE : (i+1)*CHUNK_SIZE]
		}
		f.chunks[i].id = i + 1
	}

	return
}

// Return next non-ACKed chunk 
func (f *TftpFile) GetNextChunk() (c *TftpChunk, ok bool) {
	for i := 0; i < len(f.chunks); i++ {
		if f.chunks[i].ack == false {
			return &f.chunks[i], true
		}
	}
	return nil, false
}

// Mark chink as sent
func (f *TftpFile) AckChunk(i int) {
	f.chunks[i-1].ack = true
}

func getString(buf []byte) (s string, length int) {
	for i := 0; i < len(buf); i++ {
		if buf[i] == 0 {
			return string(buf[0:i]), i
		}
	}
	return
}

func main() {
	// resolve local address
	laddr, err := net.ResolveUDPAddr("udp", ":69")
	if err != nil {
		log.Fatal("failed to resolve addr: ", err)
	}
	// listen on that port
	conn, err := net.ListenUDP("udp", laddr)
	if err != nil {
		log.Fatal("failed to create listening UDP connection: ", err)
	}

	buf := make([]byte, 512)

	var f *TftpFile

	for {
		// read single UDP packet
		n, raddr, err := conn.ReadFromUDP(buf)
		if err != nil {
			log.Println("failed to receive UDP packet: ", err)
			continue
		}
		packet := buf[0:n]

		rq := (int(packet[0]) << 8) | (int(packet[1]))

		var size int
		var filename, mode string
		if rq == RRQ {
			filename, size = getString(packet[2:])
			if size == 0 {
				log.Println("failed to retrieve file name from the packet")
				continue
			}

			mode, size = getString(packet[size+3:])
			if size == 0 {
				log.Println("failed to retrieve file mode from the packet")
				continue
			}

			mode = strings.ToLower(mode)

			if mode != MODE_OCTET && mode != MODE_ASCII {
				log.Println("mode is not supported:", mode)
				//continue
			}
		}

		switch rq {
		case RRQ:
			log.Println("RRQ", mode, filename, "from", raddr.String())
			f, err = NewTftpFile(filename)
			if err != nil {
				log.Println(err)
				continue
			}

			if c, ok := f.GetNextChunk(); ok {
				conn.WriteTo(c.Packet(), raddr)
			}
		case ACK:
			id := ((int(packet[2]) << 8) + (int(packet[3])))
			log.Println("ACK", id, "from", raddr.String())
			f.AckChunk(id)
			if c, ok := f.GetNextChunk(); ok {
				conn.WriteTo(c.Packet(), raddr)
			}
		default:
			log.Println("Unsupported request ", rq, "from", raddr.String())
		}
	}
}
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.