Make ToCTCP handle any type of CTCP message, move DCC handling to separate file

pkg/irc/client.go

@@ -3,13 +3,7 @@ package irc
 import (
 	"bufio"
 	"crypto/tls"
-	"encoding/binary"
-	"fmt"
-	"io"
-	"math"
 	"net"
-	"os"
-	"path/filepath"
 	"strings"
 	"sync"
 	"time"
@@ -218,149 +212,3 @@ func (c *Client) flushChannels() {
 	}
 	c.lock.Unlock()
 }
-
-func byteRead(totalBytes uint64) []byte {
-	b := make([]byte, 8)
-	binary.BigEndian.PutUint64(b, totalBytes)
-	return b
-}
-
-func round2(source float64) float64 {
-	return math.Round(100*source) / 100
-}
-
-const (
-	_ = 1.0 << (10 * iota)
-	kibibyte
-	mebibyte
-	gibibyte
-)
-
-func humanReadableByteCount(b float64, speed bool) string {
-	unit := ""
-	value := b
-
-	switch {
-	case b >= gibibyte:
-		unit = "GiB"
-		value = value / gibibyte
-	case b >= mebibyte:
-		unit = "MiB"
-		value = value / mebibyte
-	case b >= kibibyte:
-		unit = "KiB"
-		value = value / kibibyte
-	case b > 1 || b == 0:
-		unit = "bytes"
-	case b == 1:
-		unit = "byte"
-	}
-
-	if speed {
-		unit = unit + "/s"
-	}
-
-	stringValue := strings.TrimSuffix(
-		fmt.Sprintf("%.2f", value), ".00",
-	)
-
-	return fmt.Sprintf("%s %s", stringValue, unit)
-}
-
-func (c *Client) Download(pack *CTCP) {
-	if !c.Autoget {
-		// TODO: ask user if he/she wants to download the file
-		return
-	}
-	c.Progress <- DownloadProgress{
-		PercCompletion: 0,
-		File:           pack.File,
-	}
-	file, err := os.OpenFile(filepath.Join(c.DownloadFolder, pack.File), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
-	if err != nil {
-		c.Progress <- DownloadProgress{
-			PercCompletion: -1,
-			File:           pack.File,
-			Error:          err,
-		}
-		return
-	}
-	defer file.Close()
-
-	con, err := net.Dial("tcp", fmt.Sprintf("%s:%d", pack.IP, pack.Port))
-
-	if err != nil {
-		c.Progress <- DownloadProgress{
-			PercCompletion: -1,
-			File:           pack.File,
-			Error:          err,
-		}
-		return
-	}
-
-	defer con.Close()
-
-	var avgSpeed float64
-	var prevTime int64 = -1
-	secondsElapsed := int64(0)
-	totalBytes := uint64(0)
-	buf := make([]byte, 0, 4*1024)
-	start := time.Now().UnixNano()
-	for {
-		n, err := con.Read(buf[:cap(buf)])
-		buf = buf[:n]
-		if n == 0 {
-			if err == nil {
-				continue
-			}
-			if err == io.EOF {
-				break
-			}
-		}
-
-		if _, err := file.Write(buf); err != nil {
-			c.Progress <- DownloadProgress{
-				PercCompletion: -1,
-				File:           pack.File,
-				Error:          err,
-			}
-			return
-		}
-
-		cycleBytes := uint64(len(buf))
-		totalBytes += cycleBytes
-		percentage := round2(100 * float64(totalBytes) / float64(pack.Length))
-
-		now := time.Now().UnixNano()
-		secondsElapsed = (now - start) / 1e9
-		avgSpeed = round2(float64(totalBytes) / (float64(secondsElapsed)))
-		speed := 0.0
-
-		if prevTime < 0 {
-			speed = avgSpeed
-		} else {
-			speed = round2(1e9 * float64(cycleBytes) / (float64(now - prevTime)))
-		}
-		secondsToGo := (float64(pack.Length) - float64(totalBytes)) / speed
-		prevTime = now
-		con.Write(byteRead(totalBytes))
-		c.Progress <- DownloadProgress{
-			InstSpeed:      humanReadableByteCount(speed, true),
-			AvgSpeed:       humanReadableByteCount(avgSpeed, true),
-			PercCompletion: percentage,
-			BytesRemaining: humanReadableByteCount(float64(pack.Length-totalBytes), false),
-			BytesCompleted: humanReadableByteCount(float64(totalBytes), false),
-			SecondsElapsed: secondsElapsed,
-			SecondsToGo:    secondsToGo,
-			File:           pack.File,
-		}
-	}
-	con.Write(byteRead(totalBytes))
-	c.Progress <- DownloadProgress{
-		AvgSpeed:       humanReadableByteCount(avgSpeed, true),
-		PercCompletion: 100,
-		BytesCompleted: humanReadableByteCount(float64(totalBytes), false),
-		SecondsElapsed: secondsElapsed,
-		File:           pack.File,
-	}
-}

pkg/irc/conn.go

@@ -223,6 +223,11 @@ func (c *Client) recv() {
 				c.setNick(msg.LastParam())
 			}
 
+		case Privmsg:
+			if ctcp := msg.ToCTCP(); ctcp != nil {
+				c.handleCTCP(ctcp)
+			}
+
 		case ReplyWelcome:
 			c.setNick(msg.Params[0])
 			c.setRegistered(true)
@@ -242,10 +247,17 @@ func (c *Client) recv() {
 
 		}
 
-		if ctcp := msg.ToCTCP(); ctcp != nil {
-			go c.Download(ctcp)
-		}
-
 		c.Messages <- msg
 	}
 }
+
+func (c *Client) handleCTCP(ctcp *CTCP) {
+	switch ctcp.Command {
+	case "DCC":
+		if strings.HasPrefix(ctcp.Params, "SEND") {
+			if dccSend := ParseDCCSend(ctcp); dccSend != nil {
+				go c.Download(dccSend)
+			}
+		}
+	}
+}

pkg/irc/dcc.go

@@ -0,0 +1,227 @@
+package irc
+
+import (
+	"encoding/binary"
+	"encoding/json"
+	"fmt"
+	"io"
+	"math"
+	"net"
+	"os"
+	"path"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"time"
+)
+
+type DCCSend struct {
+	File   string `json:"file"`
+	IP     string `json:"ip"`
+	Port   uint16 `json:"port"`
+	Length uint64 `json:"length"`
+}
+
+func ParseDCCSend(ctcp *CTCP) *DCCSend {
+	params := strings.Split(ctcp.Params, " ")
+
+	if len(params) > 4 {
+		ip, err := strconv.Atoi(params[2])
+		port, err := strconv.Atoi(params[3])
+		length, err := strconv.Atoi(params[4])
+
+		if err != nil {
+			return nil
+		}
+
+		ip3 := uint32ToIP(ip)
+
+		filename := path.Base(params[1])
+		if filename == "/" || filename == "." {
+			filename = ""
+		}
+
+		return &DCCSend{
+			File:   filename,
+			IP:     ip3,
+			Port:   uint16(port),
+			Length: uint64(length),
+		}
+	}
+
+	return nil
+}
+
+func (c *Client) Download(pack *DCCSend) {
+	if !c.Autoget {
+		// TODO: ask user if he/she wants to download the file
+		return
+	}
+	c.Progress <- DownloadProgress{
+		PercCompletion: 0,
+		File:           pack.File,
+	}
+	file, err := os.OpenFile(filepath.Join(c.DownloadFolder, pack.File), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
+	if err != nil {
+		c.Progress <- DownloadProgress{
+			PercCompletion: -1,
+			File:           pack.File,
+			Error:          err,
+		}
+		return
+	}
+	defer file.Close()
+
+	con, err := net.Dial("tcp", fmt.Sprintf("%s:%d", pack.IP, pack.Port))
+
+	if err != nil {
+		c.Progress <- DownloadProgress{
+			PercCompletion: -1,
+			File:           pack.File,
+			Error:          err,
+		}
+		return
+	}
+
+	defer con.Close()
+
+	var avgSpeed float64
+	var prevTime int64 = -1
+	secondsElapsed := int64(0)
+	totalBytes := uint64(0)
+	buf := make([]byte, 0, 4*1024)
+	start := time.Now().UnixNano()
+	for {
+		n, err := con.Read(buf[:cap(buf)])
+		buf = buf[:n]
+		if n == 0 {
+			if err == nil {
+				continue
+			}
+			if err == io.EOF {
+				break
+			}
+		}
+
+		if _, err := file.Write(buf); err != nil {
+			c.Progress <- DownloadProgress{
+				PercCompletion: -1,
+				File:           pack.File,
+				Error:          err,
+			}
+			return
+		}
+
+		cycleBytes := uint64(len(buf))
+		totalBytes += cycleBytes
+		percentage := round2(100 * float64(totalBytes) / float64(pack.Length))
+
+		now := time.Now().UnixNano()
+		secondsElapsed = (now - start) / 1e9
+		avgSpeed = round2(float64(totalBytes) / (float64(secondsElapsed)))
+		speed := 0.0
+
+		if prevTime < 0 {
+			speed = avgSpeed
+		} else {
+			speed = round2(1e9 * float64(cycleBytes) / (float64(now - prevTime)))
+		}
+		secondsToGo := (float64(pack.Length) - float64(totalBytes)) / speed
+		prevTime = now
+		con.Write(byteRead(totalBytes))
+		c.Progress <- DownloadProgress{
+			InstSpeed:      humanReadableByteCount(speed, true),
+			AvgSpeed:       humanReadableByteCount(avgSpeed, true),
+			PercCompletion: percentage,
+			BytesRemaining: humanReadableByteCount(float64(pack.Length-totalBytes), false),
+			BytesCompleted: humanReadableByteCount(float64(totalBytes), false),
+			SecondsElapsed: secondsElapsed,
+			SecondsToGo:    secondsToGo,
+			File:           pack.File,
+		}
+	}
+	con.Write(byteRead(totalBytes))
+	c.Progress <- DownloadProgress{
+		AvgSpeed:       humanReadableByteCount(avgSpeed, true),
+		PercCompletion: 100,
+		BytesCompleted: humanReadableByteCount(float64(totalBytes), false),
+		SecondsElapsed: secondsElapsed,
+		File:           pack.File,
+	}
+}
+
+type DownloadProgress struct {
+	File           string  `json:"file"`
+	Error          error   `json:"error"`
+	BytesCompleted string  `json:"bytes_completed"`
+	BytesRemaining string  `json:"bytes_remaining"`
+	PercCompletion float64 `json:"perc_completion"`
+	AvgSpeed       string  `json:"avg_speed"`
+	InstSpeed      string  `json:"speed"`
+	SecondsElapsed int64   `json:"elapsed"`
+	SecondsToGo    float64 `json:"eta"`
+}
+
+func (p DownloadProgress) ToJSON() string {
+	progress, err := json.Marshal(p)
+	if err != nil {
+		return ""
+	}
+	return string(progress)
+}
+
+func uint32ToIP(n int) string {
+	var byte1 = n & 255
+	var byte2 = ((n >> 8) & 255)
+	var byte3 = ((n >> 16) & 255)
+	var byte4 = ((n >> 24) & 255)
+	return fmt.Sprintf("%d.%d.%d.%d", byte4, byte3, byte2, byte1)
+}
+
+func byteRead(totalBytes uint64) []byte {
+	b := make([]byte, 8)
+	binary.BigEndian.PutUint64(b, totalBytes)
+	return b
+}
+
+func round2(source float64) float64 {
+	return math.Round(100*source) / 100
+}
+
+const (
+	_ = 1.0 << (10 * iota)
+	kibibyte
+	mebibyte
+	gibibyte
+)
+
+func humanReadableByteCount(b float64, speed bool) string {
+	unit := ""
+	value := b
+
+	switch {
+	case b >= gibibyte:
+		unit = "GiB"
+		value = value / gibibyte
+	case b >= mebibyte:
+		unit = "MiB"
+		value = value / mebibyte
+	case b >= kibibyte:
+		unit = "KiB"
+		value = value / kibibyte
+	case b > 1 || b == 0:
+		unit = "bytes"
+	case b == 1:
+		unit = "byte"
+	}
+
+	if speed {
+		unit = unit + "/s"
+	}
+
+	stringValue := strings.TrimSuffix(
+		fmt.Sprintf("%.2f", value), ".00",
+	)
+
+	return fmt.Sprintf("%s %s", stringValue, unit)
+}

pkg/irc/message.go

@@ -1,12 +1,7 @@
 package irc
 
 import (
-	"encoding/json"
-	"fmt"
-	"path"
-	"strconv"
 	"strings"
-	"unicode"
 )
 
 type Message struct {
@@ -17,26 +12,6 @@ type Message struct {
 	Params  []string
 }
 
-type DownloadProgress struct {
-	File           string  `json:"file"`
-	Error          error   `json:"error"`
-	BytesCompleted string  `json:"bytes_completed"`
-	BytesRemaining string  `json:"bytes_remaining"`
-	PercCompletion float64 `json:"perc_completion"`
-	AvgSpeed       string  `json:"avg_speed"`
-	InstSpeed      string  `json:"speed"`
-	SecondsElapsed int64   `json:"elapsed"`
-	SecondsToGo    float64 `json:"eta"`
-}
-
-// CTCP is used to parse a message into a CTCP message
-type CTCP struct {
-	File   string `json:"file"`
-	IP     string `json:"ip"`
-	Port   uint16 `json:"port"`
-	Length uint64 `json:"length"`
-}
-
 func (m *Message) LastParam() string {
 	if len(m.Params) > 0 {
 		return m.Params[len(m.Params)-1]
@@ -44,37 +19,31 @@ func (m *Message) LastParam() string {
 	return ""
 }
 
-// ToCTCP tries to parse the message parameters into a CTCP message
+type CTCP struct {
-func (m *Message) ToCTCP() *CTCP {
+	Command string
-	params := strings.Join(m.Params, " ")
+	Params  string
-	if strings.Contains(params, "DCC SEND") {
+}
-		// to be extra sure that there are non-printable characters
-		params = strings.TrimFunc(params, func(r rune) bool {
-			return !unicode.IsPrint(r)
-		})
-		parts := strings.Split(params, " ")
-		ip, err := strconv.Atoi(parts[4])
-		port, err := strconv.Atoi(parts[5])
-		length, err := strconv.Atoi(parts[6])
 
-		if err != nil {
+func (m *Message) ToCTCP() *CTCP {
+	lp := m.LastParam()
+
+	if len(lp) > 1 && lp[0] == 0x01 {
+		parts := strings.SplitN(strings.Trim(lp, "\x01"), " ", 2)
+		ctcp := CTCP{}
+
+		if parts[0] != "" {
+			ctcp.Command = parts[0]
+		} else {
 			return nil
 		}
 
-		ip3 := uint32ToIP(ip)
+		if len(parts) == 2 {
-
+			ctcp.Params = parts[1]
-		filename := path.Base(parts[3])
-		if filename == "/" || filename == "." {
-			filename = ""
 		}
 
-		return &CTCP{
+		return &ctcp
-			File:   filename,
-			IP:     ip3,
-			Port:   uint16(port),
-			Length: uint64(length),
-		}
 	}
+
 	return nil
 }
 
@@ -171,19 +140,3 @@ var unescapeTagReplacer = strings.NewReplacer(
 func unescapeTag(s string) string {
 	return unescapeTagReplacer.Replace(s)
 }
-
-func uint32ToIP(n int) string {
-	var byte1 = n & 255
-	var byte2 = ((n >> 8) & 255)
-	var byte3 = ((n >> 16) & 255)
-	var byte4 = ((n >> 24) & 255)
-	return fmt.Sprintf("%d.%d.%d.%d", byte4, byte3, byte2, byte1)
-}
-
-func (p DownloadProgress) ToJSON() string {
-	progress, err := json.Marshal(p)
-	if err != nil {
-		return ""
-	}
-	return string(progress)
-}