github.com/annwntech/go-micro/v2@v2.9.5/agent/input/discord/conn.go (about)

     1  package discord
     2  
     3  import (
     4  	"errors"
     5  	"strings"
     6  	"sync"
     7  
     8  	"github.com/bwmarrin/discordgo"
     9  	"github.com/annwntech/go-micro/v2/agent/input"
    10  	"github.com/annwntech/go-micro/v2/logger"
    11  )
    12  
    13  type discordConn struct {
    14  	master *discordInput
    15  	exit   chan struct{}
    16  	recv   chan *discordgo.Message
    17  
    18  	sync.Mutex
    19  }
    20  
    21  func newConn(master *discordInput) *discordConn {
    22  	conn := &discordConn{
    23  		master: master,
    24  		exit:   make(chan struct{}),
    25  		recv:   make(chan *discordgo.Message),
    26  	}
    27  
    28  	conn.master.session.AddHandler(func(s *discordgo.Session, m *discordgo.MessageCreate) {
    29  		if m.Author.ID == master.botID {
    30  			return
    31  		}
    32  
    33  		whitelisted := false
    34  		for _, ID := range conn.master.whitelist {
    35  			if m.Author.ID == ID {
    36  				whitelisted = true
    37  				break
    38  			}
    39  		}
    40  
    41  		if len(master.whitelist) > 0 && !whitelisted {
    42  			return
    43  		}
    44  
    45  		var valid bool
    46  		m.Message.Content, valid = conn.master.prefixfn(m.Message.Content)
    47  		if !valid {
    48  			return
    49  		}
    50  
    51  		conn.recv <- m.Message
    52  	})
    53  
    54  	return conn
    55  }
    56  
    57  func (dc *discordConn) Recv(event *input.Event) error {
    58  	for {
    59  		select {
    60  		case <-dc.exit:
    61  			return errors.New("connection closed")
    62  		case msg := <-dc.recv:
    63  
    64  			event.From = msg.ChannelID + ":" + msg.Author.ID
    65  			event.To = dc.master.botID
    66  			event.Type = input.TextEvent
    67  			event.Data = []byte(msg.Content)
    68  			return nil
    69  		}
    70  	}
    71  }
    72  
    73  func ChunkString(s string, chunkSize int) []string {
    74  	var chunks []string
    75  	runes := []rune(s)
    76  
    77  	if len(runes) == 0 {
    78  		return []string{s}
    79  	}
    80  
    81  	for i := 0; i < len(runes); i += chunkSize {
    82  		nn := i + chunkSize
    83  		if nn > len(runes) {
    84  			nn = len(runes)
    85  		}
    86  		chunks = append(chunks, string(runes[i:nn]))
    87  	}
    88  	return chunks
    89  }
    90  
    91  func (dc *discordConn) Send(e *input.Event) error {
    92  	fields := strings.Split(e.To, ":")
    93  	for _, chunk := range ChunkString(string(e.Data), 2000) {
    94  		_, err := dc.master.session.ChannelMessageSend(fields[0], chunk)
    95  		if err != nil {
    96  			if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
    97  				logger.Error("[bot][loop][send]", err)
    98  			}
    99  		}
   100  	}
   101  	return nil
   102  }
   103  
   104  func (dc *discordConn) Close() error {
   105  	if err := dc.master.session.Close(); err != nil {
   106  		return err
   107  	}
   108  
   109  	select {
   110  	case <-dc.exit:
   111  		return nil
   112  	default:
   113  		close(dc.exit)
   114  	}
   115  	return nil
   116  }