Here is Prob. 9, Sec. 3.3, in the book Topics in Algebra by I.N. Herstein, 2nd edition:
(a) If the group $G$ has three elements, show it must be abelian.
(b) Do part (a) if $G$ has four elements.
(c) Do part (a) if $G$ has five elements.
I think I am clear on how to tackle Part (a).
So here I will only be attempting Part (b).
My Attempt:
Suppose our group $G$ has four distinct elements, say, $e, a, b, c$, with $e$ being the identity element.
Now suppose, if possible, that $ab \neq ba$.
As the elements $e, a, b, c$ of $G$ are all distinct, so by virtue of the cancellation laws (i.e. Lemma 2.3.2 in Herstein), we cannot have $ab = a$, $ab = b$, $ba = a$, or $ba = b$.
Therefore we must $ab, ba \in \{ e, c \}$.
Since $ab \neq ba$ according to our supposition, there we can assume without any loss of generality that $ab = c$ and $ba = e$.
Thus our group $G$ has the four distinct elements $e, a, b, ab$.
Since $ba = e$, we also have $$ a = ae = a(ba). \tag{1} $$
However, since $G$ is a group, we must have $$ a(ba) = (ab)a. \tag{2} $$
From (1) and (2) above, we obtain $$ a = (ab)a, $$ from which we obtain $$ e = ab, $$ again by Lemma 2.3.2 in Herstein. This contradicts the fact that $e$ and $ab$ are two distinct elements of $G$. So our supposition that $ab \neq ba$ is wrong. Therefore we must have $$ ab = ba. $$
An analogous argument yield $bc = cb$ and also $ca = ac$.
Thus any two of the elements $a, b, c$ of $G$ commute. And, the identity element $e$ of course commutes with itself as well as with each of $a$, $b$, and $c$.
Hence our group $G$ must be abelian.
Is this proof correct? If so, is it rigorous enough for Herstein? Or, are there problems and issues?
Yes your proof in fact is correct. But here are a few possible shorthands. First you only need to work with one non commuting pair of elements:
Suppose $G$ is not abelian, then there is some pair $a,b\in G$ such that: $ab\ne ba$. (We then can conclude that $a,b\ne e$). As you then found out: $ab,ba\in\{e,c\}$ for the fourth element $c\ne e,a,b$. So either $ab=e$ or $ab=c$ resulting in $ba=e$ (otherwise $a$ and $b$ commute and we are done). So in particular $a$ and $b$ are inverse elements to one another and hence commute. A contradiction. $G$ is abelian!
But you proof is just fine! :) Also nice effort of writing it down!