Refactor handler matching and validate command entity bounds

handlers.go

@@ -45,64 +45,103 @@ func (h handler) match(update *models.Update) bool {
 		return h.matchFunc(update)
 	}
 
-	var data string
-	var entities []models.MessageEntity
+	data, entities, ok := getDataFromUpdate(update, h.handlerType)
+	if !ok {
+		return false
+	}
+
+	switch h.matchType {
+	case MatchTypeExact:
+		return h.matchExact(data)
+	case MatchTypePrefix:
+		return h.matchPrefix(data)
+	case MatchTypeContains:
+		return h.matchContains(data)
+	case MatchTypeCommand:
+		return h.matchCommand(data, entities)
+	case MatchTypeCommandStartOnly:
+		return h.matchCommandStartOnly(data, entities)
+	case matchTypeRegexp:
+		return h.matchRegexp(data)
+	default:
+		return false
+	}
+}
 
-	switch h.handlerType {
+func getDataFromUpdate(update *models.Update, handlerType HandlerType) (data string, entities []models.MessageEntity, ok bool) {
+	switch handlerType {
 	case HandlerTypeMessageText:
 		if update.Message == nil {
-			return false
+			return "", nil, false
 		}
 		data = update.Message.Text
 		entities = update.Message.Entities
 	case HandlerTypeCallbackQueryData:
 		if update.CallbackQuery == nil {
-			return false
+			return "", nil, false
 		}
 		data = update.CallbackQuery.Data
 	case HandlerTypeCallbackQueryGameShortName:
 		if update.CallbackQuery == nil {
-			return false
+			return "", nil, false
 		}
 		data = update.CallbackQuery.GameShortName
 	case HandlerTypePhotoCaption:
 		if update.Message == nil {
-			return false
+			return "", nil, false
 		}
 		data = update.Message.Caption
 		entities = update.Message.CaptionEntities
 	}
+	return data, entities, true
+}
 
-	if h.matchType == MatchTypeExact {
-		return data == h.pattern
-	}
-	if h.matchType == MatchTypePrefix {
-		return strings.HasPrefix(data, h.pattern)
-	}
-	if h.matchType == MatchTypeContains {
-		return strings.Contains(data, h.pattern)
+func (h handler) matchExact(data string) bool {
+	return data == h.pattern
+}
+
+func (h handler) matchPrefix(data string) bool {
+	return strings.HasPrefix(data, h.pattern)
+}
+
+func (h handler) matchContains(data string) bool {
+	return strings.Contains(data, h.pattern)
+}
+
+func (h handler) matchRegexp(data string) bool {
+	return h.re.Match([]byte(data))
+}
+
+func extractCommand(data string, entity models.MessageEntity) (string, bool) {
+	// Check the offset before subtracting, avoiding overflow in offset + length.
+	if entity.Offset < 0 || entity.Offset > len(data) || entity.Length <= 1 || entity.Length > len(data)-entity.Offset {
+		return "", false
 	}
-	if h.matchType == MatchTypeCommand {
-		for _, e := range entities {
-			if e.Type == models.MessageEntityTypeBotCommand {
-				if data[e.Offset+1:e.Offset+e.Length] == h.pattern {
-					return true
-				}
+	// Skipping the "/" character at the beginning of the command
+	return data[entity.Offset+1 : entity.Offset+entity.Length], true
+}
+
+func (h handler) matchCommand(data string, entities []models.MessageEntity) bool {
+	for _, e := range entities {
+		if e.Type == models.MessageEntityTypeBotCommand {
+			command, ok := extractCommand(data, e)
+			if ok && command == h.pattern {
+				return true
 			}
 		}
 	}
-	if h.matchType == MatchTypeCommandStartOnly {
-		for _, e := range entities {
-			if e.Type == models.MessageEntityTypeBotCommand {
-				if e.Offset == 0 && data[e.Offset+1:e.Offset+e.Length] == h.pattern {
-					return true
-				}
+	return false
+}
+
+func (h handler) matchCommandStartOnly(data string, entities []models.MessageEntity) bool {
+	for _, e := range entities {
+		if e.Type == models.MessageEntityTypeBotCommand && e.Offset == 0 {
+			command, ok := extractCommand(data, e)
+			if ok && command == h.pattern {
+				return true
 			}
 		}
 	}
-	if h.matchType == matchTypeRegexp {
-		return h.re.Match([]byte(data))
-	}
 	return false
 }