SOCRATES: Consider again:—Where there is an agent, must there not also be a patient?
SOCRATES: And will not the patient suffer that which the agent does, and will not the suffering have the quality of the action? I mean, for example, that if a man strikes, there must be something which is stricken?
If to be punished is to suffer, then there is an agent acting on the punished; the punisher.
The punisher who punishes rightly must consequently punish justly, and then too is acting justly. The punished, the patient who the punisher is acting on justly, by Socrates’ earlier logic of cause and effect; that to burn necessitates something which is burned, is thus suffering justly. Previously, Socrates had his interlocutor acquiesce that what is just is honorable. The honorable is described as “either pleasant or useful”, which is likened to be the description of that which is good. And so, if the punished is suffering justly, they are suffering honorably and are suffering from what is good--benefitting.
EDIT:
I dont understand why the suffering should have the quality of the action.
In response to this, I'll lift more quotes from Plato's Gorgias which further illustrate the relationship Plato is trying to illuminate between a cause and an effect sharing the same quality.
SOCRATES: And if the striker strikes violently or quickly, that which is struck will be struck violently or quickly?
SOCRATES: And the suffering to him who is stricken is of the same nature as the act of him who strikes?
SOCRATES: And if a man burns, there is something which is burned?
SOCRATES: And if he burns in excess or so as to cause pain, the thing burned will be burned in the same way?
If something cuts, there is something being cut. With that, if something cuts well, then theres something being cut well.
Maybe don't think of the action as being inflicting suffering, but rather that the agent in this case, the punisher, is inflicting just punishment. Then, the patient; the punished, is suffering just punishment.
On top of this, what is meant by "just"?
Fair question, however, if you're familiar with Plato's dialogues, you'll know that seeking what is meant by "just" is the very topic of some of his dialogues. Meno being one of them, which ends without being able to define what the just is...
I thought I had earlier mentioned that this idea of punishment being desirable is from Gorgias. This site has an introduction which talks about the concept, and I'd really recommend Gorgias as a whole.